目录
torch.argmax
torch.argmax(input, dim=None, keepdim=False)
torch.Tensor.argmax(dim=None, keepdim=False)
用于从softmax结果中选择最大值的下标
Example:
>>> a = torch.randn(4, 4)
>>> a
tensor([[ 1.3398, 0.2663, -0.2686, 0.2450],
[-0.7401, -0.8805, -0.3402, -1.1936],
[ 0.4907, -1.3948, -1.0691, -0.3132],
[-1.6092, 0.5419, -0.2993, 0.3195]])
>>> torch.argmax(a, dim=1)
tensor([ 0, 2, 0, 1])