Thanks for tips! Now I actually have a doubt with the code i’m using regarding how the masks are loaded. Since I have four classes in the masks, using the command PIL.nearest I get basically a 1x64x64 image while if i’m not wrong is a tensor like 1xnb_classx64x64 or i’m wrong?
I was looking here pytorch_unet 18 at the line:
import torchvision.utils

def reverse_transform(inp):
    inp = inp.numpy().transpose((1, 2, 0))
    inp = np.clip(inp, 0, 1)
    inp = (inp * 255).astype(np.uint8)
    
    return inp

# Get a batch of training data
inputs, masks = next(iter(dataloaders['train']))

print(inputs.shape, masks.shape)
for x in [inputs.numpy(), masks.numpy()]:
    print(x.min(), x.max(), x.mean(), x.std())

plt.imshow(reverse_transform(inputs[3]))

where basically the outpur is 25x6x192x192. So I guess the way I’m loading the masks is wrong. Any idea?