Hi,
the problem is that argmax is a not a differentiable operation (it is piecewise constant). So no gradient can flow back through the argmax operation.
And so the loss you get at the end does not require gradients because it is not computed from any input that requires gradient in a differentiable manner.
You most likely want to use softmax here and a criterion for classification that can handle that (like nll_loss).