I assume you are working on a binary classification.
If that’s the case, note that nn.Softmax would output a tensor of ones for an input tensor (activation coming from the previous linear layer) in the shape [batch_size, 1], so you might want the use nn.Sigmoid instead.
For better numerical stability I would recommend to remove the activation (no softmax or sigmoid) and use nn.BCEWithLogitsLoss as the criterion.