Merge commit '98ac4542e0e097cd1b26c62d0ffe7fb37230347c'

MeJn · Aug 16, 2017 · 2579b6b · 2579b6b
2 parents 0d34a64 + 98ac454
commit 2579b6b
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/torch/lib/TH/generic/THTensorMath.c b/torch/lib/TH/generic/THTensorMath.c
@@ -1595,6 +1595,10 @@ void THTensor_(max)(THTensor *values_, THLongTensor *indices_, THTensor *t, int
     THLongTensor_zero(indices_);
 
     if(t->size[dimension] == 1) {
+      if (!keepdim) {
+        THTensor_(squeeze1d)(values_, values_, dimension);
+        THLongTensor_squeeze1d(indices_, indices_, dimension);
+      }
       return;
     }
 
@@ -1671,6 +1675,10 @@ void THTensor_(min)(THTensor *values_, THLongTensor *indices_, THTensor *t, int
     THLongTensor_zero(indices_);
 
     if(t->size[dimension] == 1) {
+      if (!keepdim) {
+        THTensor_(squeeze1d)(values_, values_, dimension);
+        THLongTensor_squeeze1d(indices_, indices_, dimension);
+      }
       return;
     }