From 98ac4542e0e097cd1b26c62d0ffe7fb37230347c Mon Sep 17 00:00:00 2001 From: Anton Osokin Date: Mon, 7 Aug 2017 15:33:07 +0200 Subject: [PATCH] fixing the bug with squeezing a singleton dimension in torch.min and torch.max --- generic/THTensorMath.c | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/generic/THTensorMath.c b/generic/THTensorMath.c index 2ba394c4092c3c..9d2a7b4b5db1a4 100644 --- a/generic/THTensorMath.c +++ b/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; }