Thank you for your response. That makes perfect sense.
Unfortunately after training using 1 training image and 1 testing image, the model performance is incorrect. I know I need more data but I at least expect to see some colors (green or pink).
The model predicted the following mask to be all white:

image1000×1000 9.8 KB

As follows:

image1000×1000 217 Bytes

With your help from the previous post, I have the following code in which I also saved the segments as images so that way I can see what each class prediction looks like.
    model.eval()
    for idx, (x, y) in enumerate(loader):
        x = x.to(device=device)
        with torch.no_grad():
            pred = torch.sigmoid(model(x))
            out = (pred > 0.5).float()
            class_to_color = [torch.tensor([10, 133, 1], device='cuda'), torch.tensor([243, 5, 247], device='cuda'), torch.tensor([255, 255, 255], device='cuda')]
            output = torch.zeros(1, 3, out.size(-2), out.size(-1), dtype=torch.float, device='cuda')
            for class_idx, color in enumerate(class_to_color):
                mask = out[:,class_idx,:,:] == torch.max(out, dim=1)[0]
                mask = mask.unsqueeze(1) # should have shape 1, 1, 100, 100
                curr_color = color.reshape(1, 3, 1, 1)
                segment = mask*curr_color # should have shape 1, 3, 100, 100
                torchvision.utils.save_image(segment.float(), f"{folder}/segment_{class_idx}.png")
                output += segment
        torchvision.utils.save_image(output, f"{folder}/pred_{idx}.png")
        torchvision.utils.save_image(y.float(), f"{folder}{idx}.png")

But each of the saved segments is a white PNG. I think even with 1 training and testing image, I should still see some pink or green (the 2 other classes) in the resulting prediction.