@ptrblck Does the output need to be passed through a softmax before this pred = torch.argmax(output, 1) ? I was doing segmentation with 4 classes as well, my model doesn’t have a softmax at the final layer and I use cross entropy loss