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