如何动态索引Pythorch中的张量？

例如，我得到一个张量：

tensor = torch.rand(12, 512, 768)

我有一个索引列表，假设它是：

[0,2,3,400,5,32,7,8,321,107,100,511]

我希望在给定索引列表的维度2中从512个元素中选择一个元素。然后张量的大小将变成（12，1768）。在

有办法吗？在

还有一种方法就是使用PyTorch并使用索引和torch.split公司名称：

tensor = torch.rand(12, 512, 768)

# create tensor with idx
idx_list = [0,2,3,400,5,32,7,8,321,107,100,511]
# convert list to tensor
idx_tensor = torch.tensor(idx_list) 

# indexing and splitting
list_of_tensors = tensor[:, idx_tensor, :].split(1, dim=1)

当您调用张量[：，idx_tensor，：]时，您将得到形状为：
的张量（12，len_of_idx_list，768）。
第二个维度取决于索引的数量

使用torch.split这个张量被分割成一个形状张量的列表：（12，1768）

最后，list_张量包含形状的张量：

[torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768]),
 torch.Size([12, 1, 768])]