Actually sorry the issue has not been resolved.
For a UNET multi-class segmentation model, the output should be [batch_size, # of classes, height, width].
So with 5 classes, a batch size of 1 and an image size of 100x100, the output should be [1, 5, 100, 100] which is correct.
The issue is trying to save this prediction directly as an image because PIL only recognizes 1 or 3 channels which makes sense. How can I save my model output as an image if the shape is technically incorrect for an image? Thank you.