Индексирование Pytorch

Question

У меня есть тензор [[1,2], [4,5], [7,8]] и тензор с индексами [0,1,0].

Я хочу их применитьво второе измерение, чтобы оно вернулось: [1,5,8].

Как мне это сделать?

Спасибо!

Answer 1

import torch

arr=torch.tensor([[1,2],[4,5],[7,8]])
indices_arr=torch.tensor([0,1,0])

ret=arr[[0,1,2],indices_arr]
# print(ret)
# tensor([1, 5, 7])

Answer 2

Предполагая, что вы имеете в виду [1, 5, 7] для вывода:

Одним из решений было бы объединение тензора всех индексов для измерения 0 и желаемых индексов для измерения 1.

tensor = torch.tensor([[1,2],[4,5],[7,8]])
indices = torch.tensor([0,1,0])
output = tensor[torch.arange(0, tensor.size[0]), indices]

вывод:

tensor([1, 5, 7])