Pythorch-选择softmax层后的最佳概率

我有一个使用Pytorch 0.4.0的logistic回归模型，输入是高维的，输出必须是标量-0，1或2

我使用了一个线性层和一个softmax层来返回一个nx3张量，其中每一列表示输入落入三个类之一的概率（0，1或2）

但是，我必须返回一个nx1张量，因此我需要为每个输入选择最高概率，并创建一个张量，指示哪个类具有最高的概率。如何使用Pythorch实现这一点？在

为了说明这一点，我的Softmax输出如下：

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

我必须归还这个：

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

在torch.argmax()你可能想要的是：

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