Pytorch-UNet
Pytorch-UNet copied to clipboard
Prediction not running for 5 class semantic segmentation
in the prediction.py file, the probs dimension in my implementation is (5,256,256) which incites the error of channels>4 not allowed. I get it that its the .toPILImage function as it accepts a channel range of <4. But my question is that how can we run this code for a class problem which is greater then 4 then, when we would get the above error everytime?