oneflow.select

oneflow.select()

Slices the self tensor along the selected dimension at the given index. This function returns a view of the original tensor with the given dimension removed.

Parameters

  • input (Tensor) – the input tensor.

  • dim (int) – the dimension to slice.

  • select (int) – the index to select with.

Returns

the output Tensor.

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow
>>> input = flow.rand(3, 4, 5)
>>> out = flow.select(input, 0, 1)
>>> out.size()
oneflow.Size([4, 5])
>>> out = flow.select(input, 1, 1)
>>> out.size()
oneflow.Size([3, 5])