Hi Anshu!
The short answer is use CrossEntropyLoss.



 Anshu_Garg:

I will get model output which has dimension (5,10)


This is fine.  You want the input to CrossEntropyLoss (the output
of your model) to have shape [nBatch = 5, nClass = 10].


Data Loader will provide us the labels list which size will be (5,1).


I don’t know what DataLoader provides, but you want the target
for CrossEntropyLoss (your labels) to have shape [nBatch = 5].
If you get labels with shape [nBatch, 1], squeeze() it to get rid of
the singleton dimension.


But ‘original_labels’ - contains actual labels (for example 5, 8, 9, 3, 0  etc)


This is what you want.  CrossEntropyLoss expects a target that is a
LongTensor of integer class labels that range from 0 to nClass - 1.


Or does this work is internally done by nn.MSELoss function?


Just to be clear, MSELoss is not a very good loss function for
classification.  Use CrossEntropyLoss.
Best.
K. Frank