alicanakca:

mask’s shape is torch.Size([256, 256]).


This is the issue – the mask is 2-dimensional, but you’ve provided 3 arguments to mask.permute().
I am guessing that you’re converting the image from h x w x c format to c x h x w. However, looks like the mask is only in an h x w format.