torch.argmax is not differentiable and would thus detach the output from the computation graph.
This should also yield an error such as:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

so I’m not sure if you are rewrapping the loss in a new tensor or why it’s not raised.
I’m not sure where this implementation comes from, but note that other implementations 9 are using softmax to calculate the probabilities for each class.