Prediction not running for 5 class semantic segmentation

[mjehanzaib999]

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?