torch.argmax()函数-截至2023年8月28日

argmax函数参数dim=0表示从列获取最大值索引，dim=1从行获取最大值索引，dim=-1从最后一个维度获取最大值索引^[1]。

举例

import torch
d = torch.tensor([[9,7,6],
				[4,8,2],
                 [5,10,0]])
print(torch.argmax(d , dim=0))#结果应为9，10，6的所在列的索引==》0，2，0
print(torch.argmax(d , dim=1))#结果应为9，8，10所在行的索引==》0，1，1
print(torch.argmax(d , dim=-1))#结果应为9，8，10所在行的索引==》0，1，1

运行结果