How can we get the index of the maximum value in a CUDA Tensor?

How can we get the index of the maximum value in a CUDA Tensor?
Prediction: tensor([-0.0640, -3.3598, 1.3133, -0.6903, 2.2058, -1.5199, 0.0641],
device=‘cuda:0’, grad_fn=)

1 Like

i believe you used log_softmax fn, try this
probability, index = torch.max(torch.exp(prediction), dim=1)
this will give you the highest probability with its index(label).