Do you get an error or what is now working in your code?
If your model is just not training, it’s most likely due to the last softmax layer in your model.
It looks like you are dealing with a classification use case, i.e. you are most likely using nn.CrossEntropyLoss or nn.NLLLoss.
In the former case you should just pass the logits to the criterion (i.e. no non-linearity at the end of your model), while nn.NLLLoss needs nn.LogSoftmax() as the last layer.