Hello,
Your loss seems correct, have you different classes on your dataset?
Using argmax is probably not a good idea since it is not differentiable, you maybe should try to encode your labels in a one hot vector to overcome that problem.