Pytorch - Pick Best Probability After Softmax Layer

I have a logistic regression model using Pytorch 0.4.0, where my input is high-dimensional and my output must be a scalar - 0, 1 or 2. I'm using a linear layer combined with a soft

Solution 1:

torch.argmax() is probably what you want:

import torch

x = torch.FloatTensor([[0.2, 0.1, 0.7],
                       [0.6, 0.2, 0.2],
                       [0.1, 0.8, 0.1]])

y = torch.argmax(x, dim=1)
print(y.detach())
# tensor([ 2,  0,  1])

# If you want to reshape:
y = y.view(1, -1)
print(y.detach())
# tensor([[ 2,  0,  1]])