Skip to content

Commit 51385b3

Browse files
committed
Merge commit '73e0b3f4014b9f5b716eb1216d11f13347207f27'
2 parents 2579b6b + 73e0b3f commit 51385b3

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

torch/lib/THC/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)