Pytorch - Pick best probability after softmax layer

Question

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 softmax layer to return a n x 3 tensor, where each column represents the probability of the input falling in one of the three classes (0, 1 or 2).

However, I must return a n x 1 tensor, so I need to somehow pick the highest probability for each input and create a tensor indicating which class had the highest probability. How can I achieve this using Pytorch?

To illustrate, my Softmax outputs this:

[[0.2, 0.1, 0.7],
 [0.6, 0.2, 0.2],
 [0.1, 0.8, 0.1]]

And I must return this:

[[2],
 [0],
 [1]]

Answer 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]])