Solved by vmirly1 in post #2 
              
              
                You can use .permute to swap axes and then apply .view to merge the last two dimensions. 
>>> d = torch.randn(10, 3, 105, 1024)
>>> d.shape
torch.Size([10, 3, 105, 1024])
>>> d = d.permute(0, 2, 1, 3)
>>> d.shape
torch.Size([10, 105, 3, 1024])
>>> d = d.contiguous().view(10, 105, -1)
>>> d.shape
tor…