I have 5 classes. I’m using a softmax function and getting 5 probabilities in each row that add up to 1 in total.

To calculate the loss do I just pick the column with the highest probability, assign it to the column number and then calculate loss. For eg: if column 1 has highest probability- that row is classified as Class 1 and let’s say the true output is Class 0. So the MSE loss is (1-0)^2 = 1. Or is there any better method?

Hi jpj!

Yes, you should use CrossEntropyLoss as your loss criterions (and not

use `Softmax`

because `CrossEntropyLoss`

expects *logits* rather than

probabilities as its predicted `input`

).

Best.

K. Frank

1 Like