Hi,
My multi-class UNET model output is of the following shape: [1, 6, 100, 100] which is expected because the batch size is 1, I have 6 classes, and the image size is 100x100.
How can I save a prediction as an image which contains all 6 classes using torchvision.utils. save_image? The tensor in save_image expects either 1 or 3 channels but my model technically outputs 6 channels for the 6 classes.
Thank you.