Solved by thomasjo in post #2 
              
              
                I think the confusion here stems from the fact that PyTorch by default uses a NCHW memory format, with tensor dimensions structured accordingly. In other words, tensor dimensions are (batch, channel, height, width). 
Therefore, in your example, if you change the input tensor to the assumed structure…