[FIXED] Rearrange a 5D tensor in PyTorch

Issue

I have a 5D tensor in the shape of (N,C,T,H,W). I want to rearrange it using PyTorch to the form of (N,T,HW,C). How can I do that?

Solution

Naturally you can reshape the last two dimensions of your tensor by flattening your tensor from dim=-2, this will produce a shape of (N,C,T,HW):

>>> x.flatten(-2)

Then you can permute the dimensions around:

>>> x.flatten(-2).permute(0,2,3,1)

Answered By - Ivan