   THCIndex_t *indices_data;
 
   THCUNN_assertSameGPU(state, 4, input, gradOutput, gradInput, indices);
-  THNN_(TemporalMaxPooling_shapeCheck)(state, input, kW, dW);
+  THNN_(TemporalMaxPooling_shapeCheck)(state, input, gradOutput, indices, kW, dW);
   THCTensor_(resizeAs)(state, gradInput, input);
   THCTensor_(zero)(state, gradInput);
 

  THCIndex_t *indices_data;

  THCUNN_assertSameGPU(state, 3, input, output, indices);
  THNN_(TemporalMaxPooling_shapeCheck)(state, input, kW, dW);
  if (input->nDimension == 3)
  {
    dimT = 1;

static inline void THNN_(TemporalMaxPooling_shapeCheck)(
                         THCState *state,
                         THCTensor *input,
                         THCTensor *gradOutput,
                         THCIndexTensor *indices,
                         int kW, int dW) {
  int dimT = 0; // Temporal dimension
  int dimF = 1; // Feature dimension
  int input_w;
  int input_n;
  int output_w;
  int ndims = input->nDimension;

  if (ndims == 3)
  {
    dimT = 1;
    dimF = 2;
  }
  THArgCheck(kW > 0, 5,
             "kernel size should be greater than zero, but got kW: %d", kW);
  THArgCheck(dW > 0, 6,
