返回最大值的index pytorch

返回最大值的index

import torch
a=torch.tensor([[.1,.2,.3],
                [1.1,1.2,1.3],
                [2.1,2.2,2.3],
                [3.1,3.2,3.3]])
print(a.argmax(dim=1))
print(a.argmax())

输出：

tensor([ 2,  2,  2,  2])
tensor(11)