Your code looks generally good. 
Some small suggestions:

Don’t use the .data attribute, as it could yield unwanted side effects
you can get the predictions directly via torch.argmax(output, dim=1) without applying the softmax, since the max. logit will also have the max. probability. However, you can of course apply the softmax, if you want to see the probabilities (just don’t pass it to the criterion to calculate the loss).