     gradOutput = THCudaTensor_newContiguous(state, gradOutput);
 
     int n = THCudaTensor_nElement(state, input);
-    int dim = n / nOutputPlane;
+    int mapSize = 1;
     if(ndim == 3)
-      dim /= (input->size[1] * input->size[2]);
+      mapSize = (input->size[1] * input->size[2]);
     else if(ndim == 4)
-      dim /= (input->size[2] * input->size[3]);
+      mapSize = (input->size[2] * input->size[3]);
+    int nElemsPerSample = nOutputPlane * mapSize;
     preluBackward<<<GET_BLOCKS(n), CUDA_NUM_THREADS, 0, THCState_getCurrentStream(state)>>> (
 	THCudaTensor_data(state, gradInput),
 	THCudaTensor_data(state, input),
 	w,
 	THCudaTensor_data(state, gradOutput),
-	n, nOutputPlane, dim);
+	n, nElemsPerSample, mapSize);
 
     THCudaTensor_free(state, input);
     THCudaTensor_free(state, gradOutput);

    const float *input,
    const float *weight,
    const float *gradOutput,
    int n, int nOutputPlane, int dim)
{
  CUDA_KERNEL_LOOP(i, n)
  {
    int j = (i / dim) % nOutputPlane;
    gradInput[i] = input[i] > 0 ? gradOutput[i] : gradOutput[i] * weight[j];
  }
}


};

__global__ void preluForward(float *output, const float *input, const float *weight,
    int n, int nElemsPerSample, int mapSize)
{
  CUDA_KERNEL_LOOP(i, n)
  {
    int positionInSample = i % nElemsPerSample;
    int mapNumber = positionInSample / mapSize;
    output[i] = input[i] > 0 ? input[i] : input[i] * weight[mapNumber];
  }
}

