diff options
| author | Anton Osokin <anton.osokin@gmail.com> | 2017-08-07 16:33:07 +0300 |
|---|---|---|
| committer | Soumith Chintala <soumith@gmail.com> | 2017-08-17 00:51:24 +0300 |
| commit | a0bf77ff070ca27eb2de31c6465f8ffa4e399be2 (patch) | |
| tree | fdf9301c8464425760bd0c78c4a3d3a6e73a0477 | |
| parent | 0ed5623489039ed42d3920e60070e78125743a0a (diff) | |
fixing the bug with squeezing a singleton dimension in torch.min and torch.max
| -rw-r--r-- | lib/TH/generic/THTensorMath.c | 8 |
1 files changed, 8 insertions, 0 deletions
diff --git a/lib/TH/generic/THTensorMath.c b/lib/TH/generic/THTensorMath.c index 2ba394c..9d2a7b4 100644 --- a/lib/TH/generic/THTensorMath.c +++ b/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; } |