Hey! As this thread discusses, the index returned may be randomly selected from all max values:
https://discuss.pytorch.org/t/function-torch-max-return-indices-inconsistency-between-cup-and-gpu/1890/2
Here is a link to a solution to this on stackoverflow:
https://stackoverflow.com/questions/55139801/index-selection-in-case-of-conflict-in-pytorch-argmax/55146371#55146371