Hi again
I’ve recently done something similar 
I think for the nn.CrossEntropy to work,
your mask/target should have the shape:
torch.Size([batch_size, 388, 388])
while your output from the unet should have the shape:
torch.Size([batch_size, 2, 388, 388])
I’d guess a simple fix is to change this line:
# Crop the mask to the desired output size
mask = transforms.CenterCrop(img_output_size)(mask)

so you don’t get the second channel which the output has 
BR