Взятие максимальных значений каждой строки в тензоре [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). Он вернет максимальные значения с их индексами.