Hello Blade!



 blade:

It seems that the torch.argmax function brakes the backpropagation.


argmax() is not usefully differentiable, and so, indeed, does break
backpropagation.


Does the model architecture have a problem?


Yes, the argmax() piece isn’t differentiable.
Make sure you understand why argmax() isn’t differentiable, and then
see if you can reformulate what you are doing in a way that avoids the
discrete jumps inherent in argmax().
Good luck.
K. Frank