Changing the classifier to softmax

[sickerin]

I would like to change the last layer to softmax, since getting class probabilities is important to my project. I'm mainly using the tf_efficientnet_b2_ns model, which if I understand correctly by default will use nn.Linear. If I want to use softmax is the following change in classifier.py correct?

   else:
        fc = nn.Softmax(dim=0)
        # fc = nn.Linear(num_pooled_features, num_classes, bias=True)
    return global_pool, fc

Moreover, will there be a problem if I have trained models using the default classifier (nn.Linear). But in the inference, I change the classifier to Softmax. I believe it should still give the same results, except now with class probabilities.

I'm trying them out, but just want to double-check if this makes sense.

Edit: I read a bit more, and found out that since we are using nn.CrossEntropyLoss, we should not use nn.Softmax, since it's already applying log-softmax followed by nll-loss. Then I guess this means that during training I should just use nn.Linear, but during inference, I should apply nn.Softmax.