pytorch index

• torch.index_select(input, dim, index, *, out=None)
• 该函数返回一个tensor

• Parameters

• input(Tensor), 输入的tensor
• dim(int): 需要tensor的指定维度
• index (IntTensor or LongTensor) – the 1-D tensor containing the indices to index

代码实例:

x = torch.randn(3, 4)
x
indices = torch.tensor([0, 2])
torch.index_select(x, 0, indices)
torch.index_select(x, 1, indices)

'''
tensor([[ 0.1427,  0.0231, -0.5414, -1.0009],
        [-0.4664,  0.2647, -0.1228, -1.1068],
        [-1.1734, -0.6571,  0.7230, -0.6004]])
tensor([[ 0.1427,  0.0231, -0.5414, -1.0009],
        [-1.1734, -0.6571,  0.7230, -0.6004]])
tensor([[ 0.1427, -0.5414],
        [-0.4664, -0.1228],
        [-1.1734,  0.7230]])
'''

参考文献:pytorch文档