Softmax applied twice when computing loss in LogisticRegression

[garryod]

🐛 Bug

In the LogisticRegression model, loss is computed as the negative log-likelihood of the log-softmax of the softmax of the linear classifier outputs. I.e. loss = F.nll_loss(F.log_softmax(F.softmax(self.linear(x))), y)

To Reproduce

Code sample

model = LogisticRegression(42, 42)
batch = (torch.rand(42), torch.rand(42))
model.training_step(batch, 0)

Expected behavior

Network loss should be computed as the cross entropy of predictions and targets (negative log likelihood of the log-softmax of the linear classifier outputs). I.e. loss = F.nll_loss(F.log_softmax(self.linear(x)), y)

Environment

• PyTorch Version (e.g., 1.0): 1.8.1
• OS (e.g., Linux): RHEL 8.4
• How you installed PyTorch (conda, pip, source): pip
• Build command you used (if compiling from source): N/A
• Python version: 3.9.2
• CUDA/cuDNN version: N/A
• GPU models and configuration: N/A
• Any other relevant information: None

Additional context

This bug should not cause any serious issue, however it deviates from expected behaviour and reduces value separation which has potential to negatively impact training.
The PyTorch implementation of Cross Entropy Loss includes the log-softmax, as such it expects raw unnormalized scores for each class (docs, functional)
Happy to submit a PR to resolve if deemed appropriate

[github-actions]

Hi! thanks for your contribution!, great first issue!