Trinayan_Baruah:

At the same time, we normally call backward on the loss instead of the output correct?.


Yes, that’s right.



 Trinayan_Baruah:

Does it mean that in a real network for example, we will see the backward pass of batchnorm not applying the channels last format?


No, if you add another layer (e.g. nn.Conv2d) you would see the channels-last kernels again. The issue with your code would be the reduction, which would drop the memory layout, if I’m not mistaken.