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).