If you replaced the last linear layer with model.fc = nn.Linear(2048, 1), I’m not sure how you are getting class predictions in the range [0, 4], as the model output would be a single value for each sample.
Could you explain how you get these predictions, please?



 mobassir94:

but if i use crossentropy and do loss = criterion(outputs, torch.max(labels, 1)[1]) then i get loss = 0.000 any help?


As explained before:



 ptrblck:

In your case the last linear layer has a single output unit, which would define nb_classes = 1 .
[…]
In this setup your model won’t learn anything, as it only needs to output a high value for the single output unit to predict the “right” class0.