     int lane = threadIdx.x % 32; // from 0 to 31
 
     // if less than 1024 threads per block, then only activate the relevant lanes
-    if (lane < blockDim.x / 32) 
+    if (lane < blockDim.x / 32)
     {
 #pragma unroll
       for (int i = 0; i < 32; ++i)

  THCudaTensor_free(state, gradOutput);
  THCudaTensor_free(state, output);
}

#undef LOGSOFTMAX_THREADS

{
  THCUNN_assertSameGPU(state, 3, output, gradOutput, gradInput);

  THCudaTensor_resizeAs(state, gradInput, output);

  bool spatial  = false;
  int batchSize = 1;
  int classSize = 0;
  int height = 0;
  int width = 0;

  int ndims = THCudaTensor_nDimension(state, input);

  if (ndims == 1)
  {
    classSize = THCudaTensor_size(state, gradInput, 0);
    output = THCudaTensor_newContiguous(state, output);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 2)
  {
    batchSize = THCudaTensor_size(state, gradInput, 0);
    classSize = THCudaTensor_size(state, gradInput, 1);
    output = THCudaTensor_newContiguous(state, output);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 3)
  {
    spatial = true;
    classSize = THCudaTensor_size(state, input, 0);
    height = THCudaTensor_size(state, input, 1);
    width = THCudaTensor_size(state, input, 2);

    THCudaTensor *transposedOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedOutput, output);
    THCudaTensor_copy(state, transposedOutput, output);
    output = transposedOutput;
    // convert tensor from torch layout to cuda layout
    // P x H x W -> W x H x P
    THCudaTensor_transpose(state, output, output, 0, 2);
    // W x H x P -> H x W x P
    THCudaTensor_transpose(state, output, output, 0, 1);
    output = THCudaTensor_newContiguous(state, output);

    THCudaTensor *transposedGradOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedGradOutput, gradOutput);
    THCudaTensor_copy(state, transposedGradOutput, gradOutput);
    gradOutput = transposedGradOutput;
    // convert tensor from torch layout to cuda layout
    // P x H x W -> W x H x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 2);
    // W x H x P -> H x W x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 1);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else if (ndims == 4)
  {
    spatial = true;
    batchSize = THCudaTensor_size(state, gradInput, 0);
    classSize = THCudaTensor_size(state, input, 1);
    height = THCudaTensor_size(state, input, 2);
    width = THCudaTensor_size(state, input, 3);

    THCudaTensor *transposedOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedOutput, output);
    THCudaTensor_copy(state, transposedOutput, output);
    output = transposedOutput;
    // convert tensor from torch layout to cuda layout
    // B x P x H x W -> B x W x H x P
    THCudaTensor_transpose(state, output, output, 1, 3);
    // B x W x H x P -> B x H x W x P
    THCudaTensor_transpose(state, output, output, 1, 2);
    output = THCudaTensor_newContiguous(state, output);

    THCudaTensor *transposedGradOutput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedGradOutput, gradOutput);
    THCudaTensor_copy(state, transposedGradOutput, gradOutput);
    gradOutput = transposedGradOutput;
    // convert tensor from torch layout to cuda layout
    // B x P x H x W -> B x W x H x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 3);
    // B x W x H x P -> B x H x W x P
    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 2);
    gradOutput = THCudaTensor_newContiguous(state, gradOutput);
  }
  else
  {
    THError("1D, 2D, 3D or 4D Tensor expected");
  }

  if (!spatial)
  {
    dim3 grid(batchSize);
    dim3 block(1024);
