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:41 +0300 |
| commit | 2de2d2bf67e4a4e5d139c729c26c7d609dbd6349 (patch) | |
| tree | 7559d832e9eefebb699bfabe92c564b6ece03fa9 | |
| parent | 32a9708dc2edb94ed73f6f5fa72c3b600ee598a8 (diff) | |
fixing the bug with squeezing a singleton dimension in torch.min and torch.max
| -rw-r--r-- | lib/THC/generic/THCTensorMode.cu | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/lib/THC/generic/THCTensorMode.cu b/lib/THC/generic/THCTensorMode.cu index e5a17f2..4cbcac0 100644 --- a/lib/THC/generic/THCTensorMode.cu +++ b/lib/THC/generic/THCTensorMode.cu @@ -190,6 +190,10 @@ THC_API void THCTensor_(mode)(THCState *state, if (sliceSize == 1) { THCTensor_(copy)(state, values, input); THCudaLongTensor_fill(state, indices, TH_INDEX_BASE); + if (!keepdim) { + THCTensor_(squeeze1d)(state, values, values, dimension); + THCudaLongTensor_squeeze1d(state, indices, indices, dimension); + } return; } |