Блог

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

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

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

Спасибо!

 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])
  

Предполагая, что вы имеете в виду [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])