How to reshape a 32x512 into 128x128 keeping it as consistant 4 32x128 parts stacked?
So I want to turn 32x512 into 128x128 keeping it as consistant 4 32x128 parts stacked like this:
.
How to do such a thing inside the PyTorch model?

So here is what I tried based on this answer:
import torch 
x = torch.arange(32*512*3*2).reshape(2, 3, 32, 512)
print(x.shape)

x = list(x.split(4, dim=3))
print(f&quot; xx {len(x)} {iii.shape for iii in x }&quot;)

x = torch.cat(x)
print(f&quot; xs {x.shape}&quot;)

And this outputs:
torch.Size([2, 3, 32, 512])
 xx 128 &lt;generator object &lt;genexpr&gt; at 0x7f358f1edad0&gt;
 xs torch.Size([256, 3, 32, 4])

While I wanted  xs torch.Size([2, 3, 128, 128])

You can use, for tensor t:
torch.cat(t.split(128, dim=1))

To transform in the reverse direction, you can use:
torch.cat(t.split(32, dim=0))


For your updated question:
torch.cat(t.split(128, dim=3), dim=2)

For the reverse:
torch.cat(t.split(32, dim=2), dim=3)

In general, the dim for torch.split tells which dimension you want to split over, and the dim for torch.cat tells which dimension you want to concatenate over.
