Hi,
From your code, it seems like that you are passing actual prediction value (not logits) to the loss function. It depends on which loss function you are using, nn.CrossEntropyLoss and nn.BCEWithLogitsLoss comes with internal implementation of softmax (so you have to just pass the output of last fully connected layer to loss function). nn.NLLLoss and nn.BCELoss loss function requires softmax to be applied beforehand. But any of these loss functions don’t require the following operations you’ve performed (because these operations are implemented internally).
predictions =  torch.max(predictions.data, 1)[1]
predictions = predictions.float()

Just remove softmax operation from your model.forward function. And use any of nn.CrossEntropyLoss  and  nn.BCEWithLogitsLoss loss function and remove those two operations. For more details go to Pytorch’s official doc of loss functions 2.