torch.index_select用法

torch.index_select:通过选择索引然后去得到想要的tensor,针对比较长的tensor
torch.index_select(tensor,维度，选择的index）
代码示例：

import torch
#shape为（2，2，3）
a=torch.tensor([[[1,2,3],[4,5,6]],
                [[7,8,9],[10,11,12]]])
#选择索引0和索引2的tensor                
indices=torch.tensor([0,2])
#tensor为a,维度为2，索引为0和2
b=torch.index_select(a,2,indices)
print(b)

输出;

tensor([[[ 1,  3],
         [ 4,  6]],

        [[ 7,  9],
         [10, 12]]])