Solved by AladdinPerzon in post #2 
              
              
                My thinking is that what we want is to permute the axis so we want to do 
x = torch.randn((3, 256, 256))
x = x.permute(1,2,0) # shape (256,256,3)

And using einsum is equivalent to permuting the axis. This is not what torch.view does. I think one explanation of what view does is that it will rearran…