argmax函数的使用(pytorch版)

import torch
input = torch.tensor([[6, 0],
                      [2, 9]])
print(input.argmax())
# tensor(3) （把[[6,0],[2,9]]看成[6,0,2,9]。从0索引，第3个索引的元素最大）

print(input.argmax(0))
# tensor([0, 1])
# 沿列看，第一列第0个位置（元素'6'）最大，
# 沿列看，第二列第1个位置（元素'9'）最大。依旧是从0开始索引）

print(input.argmax(1))
# tensor([0, 1])
# 沿行看，第一行第0个位置（元素'6'）最大，
# 沿行看，第二行第1个位置（元素'9'）最大。依旧是从0开始索引）

如果一个数组中最大值个数不唯一
那么返回第一个最大值的索引

input2 = torch.tensor([0,6,9,2,1,9])
print(input2.argmax())
# tensor(2) (返回第一个最大值的索引)