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