pytorch: torch.max() 使用与理解

在使用 pytorch 进行训练时，会使用使用到改行代码：
predict = torch.max(outputs.data, 1)[1]
其中 output 为模型的输出，该函数主要用来求 tensor 的最大值。

每次看到都不太理解 torch.max() 的使用，为了下次看到或者写道时不会忘记，特意详细了解其用法。
torch.max(input:tensor, dim:index)
该函数有两个输入：
inputs: tensor，第一个参数为一个张量
dim: index，第二个参数为一个整数[-2-1]，dim=0表示计算每列的最大值，dim=1表示每行的最大值
如：

import torch

output = torch.tensor([[1, 2， 3], [3, 4， 5]])
predict = torch.max(output, dim=0)
print(predict)
predict = torch.max(output, dim=1)
print(predict)

输出为：
torch.return_types.max(
values=tensor([3, 4, 5]),
indices=tensor([1, 1, 1]))

torch.return_types.max(
values=tensor([3, 5]),
indices=tensor([2, 2]))

输出返回最大值以及对应的索引值。因此，predict = torch.max(outputs.data, 1)[1] 是计算模型中每个类别的最大值并返回其索引值，即该类别的标签值。