Repository: deeplearning4j/deeplearning4j
Issue: Unable to broadcast dimension 1 due to shape mismatch exception with Deconvolution2D layers

I tried to set up a neural net using a ComputationGraph based on some research paper. This net has an encoder-decoder like structure: it uses Deconvolution2D layers for the decoding part. After putting together the net, I found that I can run the training with batch size 1, but any higher batch size value causes the net to stop with a message similar to the below:

Exception in thread "main" java.lang.IllegalArgumentException: Unable to broadcast dimension 1 due to shape mismatch. Right shape must be 1. Left array shape: [2, 16, 33, 33], right array shape: [16, 2, 33, 33]

I spent a lot of time debugging the issue to see what did I miss. I simplified the net to the smallest possible size, it now only has a Convolution, Deconvolution2D and an Output layer. I can still reproduce my original problem: with batch size 1, it works fine, otherwise it stops.
Looking at the array shapes in the error message it seems that the minibatch size and the number of channels is mixed up whenever I receive this error.

I created a GIST for this in the hope that someone could check and give me a hint on how to fix this.
https://gist.github.com/bekkon/7fd406d21d1320d08a5d2fe82e2a0331

After tracing back the calls, I found that I can seemingly fix this error by modifying the backpropGradient method in Deconvolution2D class. At the very end of the method, it does a permutation of the previously prepared outEps:
 "outEps = outEps.permute(1, 0, 2 , 3);"

If I change this to "outEps = outEps.permute(0, 1, 2 , 3);", the exception goes away, backpropagation succeeds in the preceding layer.

Please check if this is indeed an error in Deconvolution2D. If not, please help me understand what did I do wrong in my network setup to end up with this error. (The attached net in the GIST is just for reproducing the error, it is not supposed to be meaningful, for network-help I'll share the complex net in a separate GIST  )

Thanks :)

Changed files:
deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java
deeplearning4j/deeplearning4j-cuda/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java

diff --git a/deeplearning4j/deeplearning4j-cuda/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java b/deeplearning4j/deeplearning4j-cuda/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
index 22f33a7f31..5c656c807d 100644
--- a/deeplearning4j/deeplearning4j-cuda/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
+++ b/deeplearning4j/deeplearning4j-cuda/src/test/java/org/deeplearning4j/gradientcheck/CNNGradientCheckTest.java
@@ -861,7 +861,7 @@ public class CNNGradientCheckTest extends BaseDL4JTest {
     public void testDeconvolution2D() {
         int nOut = 2;
 
-        int[] minibatchSizes = new int[]{1, 3, 1, 3, 1, 3, 1, 3};
+        int[] minibatchSizes = new int[]{1, 4, 1, 4, 1, 1, 2, 1};
         int[] kernelSizes = new int[]{1, 1, 3, 3, 1, 1, 3, 3};
         int[] strides = {1, 1, 1, 1, 2, 2, 2, 2};
         int[] dilation = {1, 2, 2, 1, 1, 1, 2, 2};
@@ -895,13 +895,23 @@ public class CNNGradientCheckTest extends BaseDL4JTest {
                     .updater(new NoOp())
                     .activation(act)
                     .list()
+                    .layer(new Deconvolution2D.Builder().name("deconvolution_2D_layer-0")
+                            .cudnnAllowFallback(false)
+                            .kernelSize(1, 1)
+                            .stride(1, 1)
+                            .dilation(0, 0)
+                            .convolutionMode(cm)
+                            .nIn(inputDepth)
+                            .nOut(inputDepth)
+                            .build())
                     .layer(new Deconvolution2D.Builder().name("deconvolution_2D_layer")
                             .cudnnAllowFallback(false)
                             .kernelSize(k, k)
                             .stride(s, s)
                             .dilation(d, d)
                             .convolutionMode(cm)
-                            .nIn(inputDepth).nOut(nOut).build());
+                            .nIn(inputDepth).nOut(nOut)
+                            .build());
 
             MultiLayerConfiguration conf = b.layer(new OutputLayer.Builder(LossFunctions.LossFunction.MCXENT)
                     .activation(Activation.SOFTMAX).nOut(nOut).build())
diff --git a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java
index e86b583fef..f609d4e06e 100644
--- a/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java
+++ b/deeplearning4j/deeplearning4j-nn/src/main/java/org/deeplearning4j/nn/layers/convolution/Deconvolution2DLayer.java
@@ -132,8 +132,6 @@ public class Deconvolution2DLayer extends ConvolutionLayer {
         retGradient.setGradientFor(DeconvolutionParamInitializer.WEIGHT_KEY, weightGradView, 'c');
         weightNoiseParams.clear();
 
-        outEps = outEps.permute(1, 0, 2 , 3);
-
         return new Pair<>(retGradient, outEps);
     }