     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
   {

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

    THCudaTensor *transposedInput = THCudaTensor_new(state);
    THCudaTensor_resizeAs(state, transposedInput, input);
    THCudaTensor_copy(state, transposedInput, input);
    input = transposedInput;
    // 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, input, input, 1, 3);
    // B x W x H x P -> B x H x W x P
    THCudaTensor_transpose(state, input, input, 1, 2);
    input = THCudaTensor_newContiguous(state, input);
  }
  else
  {

    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)
  {
