index_select

dragon.vm.torch.index_select(
  input,
  dim,
  index,
  out=None
)[source]

Select the elements along the given dimension using index.

Index should be a int64 tensor:

input = torch.tensor([[1, 2, 3], [4, 5, 6]])
index = torch.tensor([1])
print(torch.index_select(input, 0, index))  # [[4, 5, 6]]

Continuous dimensions could also be specified to index:

print(torch.index_select(input, (0, 1), index))  # [2]
Parameters:
Returns:

dragon.vm.torch.Tensor – The output tensor.