Convert 5D tensor to 4D tensor in PyTorch

In PyTorch I have a 5D tensor X of dimensions B x 9 x C x H x W. I want to convert it into a 4D tensor Y with dimensions B x 9C x H x W such that concatenation happens channel wise.

To illustrate let,

a = X[1,0,:,:,:]
b = X[1,1,:,:,:]
c = X[1,2,:,:,:]
...
i = X[1,8,:,:,:]

Then in the tensor Y, a to i should be channel wise concatenated.

Answer

Try:

import torch
x = torch.rand(3, 4, 3, 2, 6)
print(x.shape)
y=x.flatten(start_dim=1, end_dim=2)
print(y.shape)

---

torch.Size([3, 4, 3, 2, 6])
torch.Size([3, 12, 2, 6])