How to use permute on this Input and Target?
I am having errors on my semantic segmentation masks with 5 classes + 1 (background).

How do I use permute to avoid this?

Target size (torch.Size([4, 1, 320, 480, 6])) must be the same as input size (torch.Size([4, 6, 320, 480]))


You can combine permute and unsqueeze:

import torch

x = torch.rand((4, 6, 320, 480))
new_x = x.permute((0,2,3,1)).unsqueeze(1)
# new_x.shape = torch.Size([4, 1, 320, 480, 6])

