Skip to content

Commit 73e0b3f

Browse files
aosokinsoumith
authored andcommitted
fixing the bug with squeezing a singleton dimension in torch.min and torch.max
1 parent ec28630 commit 73e0b3f

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

generic/THCTensorMode.cu

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -190,6 +190,10 @@ THC_API void THCTensor_(mode)(THCState *state,
190190
if (sliceSize == 1) {
191191
THCTensor_(copy)(state, values, input);
192192
THCudaLongTensor_fill(state, indices, TH_INDEX_BASE);
193+
if (!keepdim) {
194+
THCTensor_(squeeze1d)(state, values, values, dimension);
195+
THCudaLongTensor_squeeze1d(state, indices, indices, dimension);
196+
}
193197
return;
194198
}
195199

0 commit comments

Comments
 (0)