I might not fully understand the current use case.
Are you using ToPILImage on the model output to visualize it or on the target tensors?
In case it’s just to visualize the target, you might scale the values to [0, 0.33,0.66, 1]. However, don’t transform the target during training and keep the values as the class indices [0, 1, 2, 3].