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

    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];
  }
}


    gradOutput = THCudaTensor_newContiguous(state, gradOutput);

    int n = THCudaTensor_nElement(state, input);
    int mapSize = 1;
    if(ndim == 3)
      mapSize = (input->size[1] * input->size[2]);
    else if(ndim == 4)
      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, nElemsPerSample, mapSize);

    THCudaTensor_free(state, input);
    THCudaTensor_free(state, gradOutput);
