# Torchvision models dont have softmax layer？

In the torchvision models vgg code：https://github.com/pytorch/vision/blob/master/torchvision/models/vgg.py

there is not a softmax layer in this code ,but in the original paper the last layer is a softmax layer，why dose it happen？
if I use pretrained vgg model to recognize a picture，do I need to add a softmax layer in the last of the model to get each class’s probability？

4 Likes

torchvision models were trained using`nn.CrossEntropyLoss` which consists of `nn.LogSoftmax` and then `nn.NLLLoss`.
that’s why there is no softmax layer at the end.
So you only need to use `nn.CrossEntropyLoss()` in your model

9 Likes

if you want only the predicted class then you don’t need it you can only use `torch.max()`