Hi, @thunder, I don’t think torch.max is suitable to your case. I think torch.argmax is a good choice.
torch.max
torch.argmax