oneflow.argmin

oneflow.argmin()

The op computes the index with the largest value of a Tensor at specified axis.

Parameters
  • input (oneflow.Tensor) – Input Tensor

  • dim (int, optional) – dimension to be calculated. Defaults to the last dim (-1)

  • keepdim (bool optional) – whether the output tensor has dim retained or not. Ignored if dim=None.

Returns

A Tensor(dtype=int64) contains the index with the largest value of input

Return type

oneflow.Tensor

For example:

>>> import oneflow as flow

>>> input = flow.tensor([[4, 3, 1, 0, 2],
...            [5, 9, 7, 6, 8]], dtype=flow.float32)
>>> output = flow.argmin(input)
>>> output
tensor(3, dtype=oneflow.int64)
>>> output = flow.argmin(input, dim=1)
>>> output
tensor([3, 0], dtype=oneflow.int64)