@@ -11,9 +11,10 @@ void THNN_(GatedLinear_updateOutput)(
1111 THCUNN_assertSameGPU (state, 2 , input, output);
1212
1313 // size output to half of input
14- dim = dim - 1 ;
14+ dim = dim - TH_INDEX_BASE ;
1515 const long nIn = THCTensor_ (size)(state, input, dim);
16- THArgCheck (nIn % 2 == 0 , 2 , " Halving dimension must be even. Dim %d is size %ld" , dim+1 , nIn);
16+ THArgCheck (nIn % 2 == 0 , 2 , " Halving dimension must be even. Dim %d is size %ld" ,
17+ dim + TH_INDEX_BASE, nIn);
1718 const long inputSize = THCTensor_ (size)(state, input, dim) / 2 ;
1819 THLongStorage *newSizes = THCTensor_ (newSizeOf)(state, input);
1920 THLongStorage_set (newSizes, dim, inputSize);
@@ -39,9 +40,10 @@ void THNN_(GatedLinear_updateGradInput)(
3940 int dim)
4041{
4142 THCUNN_assertSameGPU (state, 2 , gradOutput, gradInput);
42- dim = dim - 1 ;
43+ dim = dim - TH_INDEX_BASE ;
4344 const long nIn = THCTensor_ (size)(state, input, dim);
44- THArgCheck (nIn % 2 == 0 , 2 , " Halving dimension must be even. Dim %d is size %ld" , dim+1 , nIn);
45+ THArgCheck (nIn % 2 == 0 , 2 , " Halving dimension must be even. Dim %d is size %ld" ,
46+ dim + TH_INDEX_BASE, nIn);
4547
4648 THCTensor_ (resizeAs)(state, gradInput, input);
4749 const long inputSize = THCTensor_ (size)(state, input, dim) / 2 ;
@@ -61,4 +63,4 @@ void THNN_(GatedLinear_updateGradInput)(
6163 THCTensor_ (free )(state, gradInputsecondHalf);
6264}
6365
64- #endif
66+ #endif
0 commit comments