pytorch index_ select

• torch.index_select(input, dim, index, *, out=None)
• This function returns a tensor

• Parameters

• input(Tensor), Input tensor
• dim(int): need tensor The specified dimension of
• index (IntTensor or LongTensor) – the 1-D tensor containing the indices to index

Code instance :

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

reference :pytorch file