     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);
+    // create contiguous tensor with cuda layout from tensor with torch layout
+    // B x C x H x W -> B x W x H x C
+    THCudaTensor_transpose(state, output, output, 1, 3);
+    // B x W x H x C -> B x H x W x C
+    THCudaTensor_transpose(state, output, output, 1, 2);
+    THCudaTensor *transposedOutput = THCudaTensor_newContiguous(state, output);
+    THCudaTensor_transpose(state, output, output, 1, 2);
+    THCudaTensor_transpose(state, output, output, 1, 3);
     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);
+
+    // create contiguous tensor with cuda layout from tensor with torch layout
+    // B x C x H x W -> B x W x H x C
+    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 3);
+    // B x W x H x C -> B x H x W x C
+    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 2);
+    THCudaTensor *transposedGradOutput = THCudaTensor_newContiguous(state, gradOutput);
+    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 2);
+    THCudaTensor_transpose(state, gradOutput, gradOutput, 1, 3);
     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
   {

    if (y >= height)
      break;

    // calculate input starting index in cuda layout (B x H x W x P)
    int inputStartIndex = 0
    + (height*width*classSize)*batchIndex
    + (width*classSize)*y
    + (classSize)*x;

    // compute a
    float a = 0;
    for (int i = 0; i < classSize; i++) {
      a += __expf(input[inputStartIndex + i]);
    }

    // invert a
    a = 1.0f / a;

    for (int i = 0; i < classSize; i++) {
      // calculate output index in torch layout (B x P x H x W)
      int outputIndex = 0
      + (classSize*height*width)*batchIndex
      + (height*width)*i
      + (width)*y
      + x;
      output[outputIndex] = logf(a * __expf(input[inputStartIndex + i]));
    }
    index += 1024;
  }
}


    height = THCudaTensor_size(state, input, 1);
    width = THCudaTensor_size(state, input, 2);

    // create contiguous tensor with cuda layout from tensor with torch layout
    // C x H x W -> W x H x C
    THCudaTensor_transpose(state, output, output, 0, 2);
    // W x H x C -> H x W x C
    THCudaTensor_transpose(state, output, output, 0, 1);
    THCudaTensor *transposedOutput = THCudaTensor_newContiguous(state, output);
    THCudaTensor_transpose(state, output, output, 0, 1);
    THCudaTensor_transpose(state, output, output, 0, 2);
    output = transposedOutput;

    // create contiguous tensor with cuda layout from tensor with torch layout
    // C x H x W -> W x H x C
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 2);
    // W x H x C -> H x W x C
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 1);
    THCudaTensor *transposedGradOutput = THCudaTensor_newContiguous(state, gradOutput);
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 1);
    THCudaTensor_transpose(state, gradOutput, gradOutput, 0, 2);
    gradOutput = transposedGradOutput;
  }
  else if (ndims == 4)
  {
