我有张量1,2],[4,5],[7,8和张量为[0,1,0]的张量。
我想将它们应用于第二维,以便它返回:[1,5,8]。
我该怎么办?
谢谢!
假设你的意思是[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])
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])