Hi Niranjan!



 Niranjan_Rao:

Do I have to take some kind of argmax on the nb_classes dimension?


Yes, this is exactly right.  The algebraically largest logit corresponds
to the most probable class.
If you have:
logits.shape == torch.Size ([batch_size, nb_classes=4, height, width])

then logits.argmax (dim = 1) will be a LongTensor of shape
[batch_size, height, width], with the nb_classes dimension
“reduced” away.
Best.
K. Frank