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 rearrange the dimensions to match, not swap them which is what permute is doing. You might want to check out this post to see a better example of what view
is doing