How I can swap 3 dimensions with each other in Pytorch?

Question

I have a a= torch.randn(28, 28, 8) and I want to swap the dimensions of the tensor and move the third dimension to the first place, first one to the second place and the second one to the third place. I used b = a.transpose(2, 0, 1) , but I received this error:

TypeError: transpose() received an invalid combination of arguments - got (int, int, int), but expected one of:
 * (name dim0, name dim1)
 * (int dim0, int dim1)

Should I use transpose several times, each time only to swap two dimensions? Is there any way that I can swap all at once?

Thanks.

Answer 1

Use permute :

b = a.permute(2, 0, 1)

Answer 2

You can use Pytorch's permute() function to swap all at once,

>>>a = torch.randn(28, 28, 8)
>>>b = a.permute(2, 0, 1)
>>>b.shape
torch.Size([8, 28, 28])