lib/dnn/core/layers.rb in ruby-dnn-0.5.1 vs lib/dnn/core/layers.rb in ruby-dnn-0.5.2

@@ -188,14 +188,15 @@
         self.new(hash[:shape])
       end
 
       def forward(x)
         @x_shape = x.shape
-        x.reshape(*@shape)
+        x.reshape(x.shape[0], *@shape)
       end
 
       def backward(dout)
-        dout.reshape(@x_shape)
+        p dout.shape, @x_shape
+        dout.reshape(*@x_shape)
       end
 
       def to_hash
         super({shape: @shape})
       end