torch.max()函数用法

Handsome_zhou

torch.max()函数第二个参数为0时按列取最大值，参数为1时按行取最大值。

import torch
a = torch.randn(3,3)
print(a)
print(torch.max(a)) #返回tensor a中的最大值
print(torch.max(a,0))#返回每一列中最大值的那个元素，并返回索引
print(torch.max(a,1))#返回每一行中最大值的那个元素，并返回索引
print(torch.max(a,0)[0])#只返回最大值的每个数
print(torch.max(a,0)[1])#只返回最大值的每个索引



运行结果：