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:

• input (dragon.vm.torch.Tensor) – The input tensor.
• dim (Union[int, Sequence[int]]) – The dimension to select.
• index (dragon.vm.torch.Tensor) – The index tensor.
• out (dragon.vm.torch.Tensor, optional) – The output tensor.

Returns:

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