如何在PyTorch一维张量中按索引列表选取元素（无循环）

你可以通过两种简洁的方式实现，完全不需要循环：

方法一：直接张量索引 #

这是最直观的方式，PyTorch支持直接用索引张量对一维张量进行元素选取：

import torch

t = torch.tensor([10., 20., 30., 40., 50., 60.])
indices = torch.tensor([0, 3, 4])
result = t[indices]
print(result)
# 输出: tensor([10., 40., 50.])

方法二：使用torch.index_select #

torch.index_select并非只适用于二维张量，一维张量同样可以使用，只需指定维度为0（一维张量唯一的维度）：

import torch

t = torch.tensor([10., 20., 30., 40., 50., 60.])
indices = torch.tensor([0, 3, 4])
result = torch.index_select(t, dim=0, index=indices)
print(result)
# 输出: tensor([10., 40., 50.])

内容的提问来源于stack exchange，提问作者PeeteKeesel