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