Neural networks, deep learning, NLP · insight-0405

Feed raw logits to CrossEntropyLoss, never softmax outputs

Situation
A multi-class model ends in nn.Softmax and its output goes straight into nn.CrossEntropyLoss.
Action
Return logits from forward, pass them to the loss, and apply softmax only in a separate predict function used for reporting, calibration and thresholds.
Reason
PyTorch's cross-entropy expects unnormalized logits and applies log-softmax internally in a numerically stable way.
Boundary
PyTorch classification heads; check the equivalent contract in other frameworks, where some losses take a from-logits flag.

Symptom or risk

Softmax is applied twice. Gradients shrink, training is slow and plateaus early, and reported probabilities are squeezed towards uniform.

If you deliberately work in log-probabilities, use log_softmax with NLLLoss, which is equivalent.

A predict wrapper adds a code path that must stay consistent with training and serving.

Exporting the training forward to serving and reading logits as probabilities; thresholds tuned on probabilities then fail silently.

PyTorch CrossEntropyLoss expects its input to contain unnormalized logits for each class.
CrossEntropyLoss (PyTorch documentation)