取张量中每一行的最大值[PyTorch]

Question

假设我有一个形式的张量

[[-5, 0, -1],
 [3, 100, 87],
 [17, -34, 2],
 [45, 1, 25]]

我想在每一行中找到最大值并返回一个 1 阶张量，如下所示：

[0,
 100,
 17,
 45]

我将如何在 PyTorch 中执行此操作？

Answer 1

您可以使用torch.max()函数。 所以你可以做类似的事情

x = torch.Tensor([[-5, 0, -1],
                  [3, 100, 87],
                  [17, -34, 2],
                  [45, 1, 25]])
out, inds = torch.max(x,dim=1)

这将返回每行（维度 1）的最大值。 它将返回最大值及其索引。