Hi,
The argmax operation is not differentiable, if you check the two outputs of the max operation, you will see that the value has requires_grad=True (this one is differentiable) but the indices has requires_grad=False (this one is not).
You cannot learn this unfortunately, to do classification, you will need a criterion that works directly with the outputs values like nn.CrossEntropyLoss