Thanks for the quick reply!
I must admit that is not clear to me how to create the mask. Once I load the mask with PIL I pass my mask as

mask = torch.from_numpy(mask)

but then is not clear how to create this one:



 ptrblck:

but instead contain the class indices in the shape [batch_size, height, width] .


Because I guess once I create this tensor like [batch_size, num_classes, height , width] I guess I can skip this step on the training loop



 Diauro:

target_ = torch.empty(batch_size, 1,64,64)  target = target_.to(device)


That gives me the error with the Nan due to what you said about the memory that might contain those values.
Best