megengine.functional.nn.one_hot

one_hot(inp, num_classes)[source]

Performs one-hot encoding for the input tensor.

Parameters
  • inp (Tensor) – input tensor.

  • num_classes (int) – number of classes denotes the last dimension of the output tensor.

Examples

>>> import numpy as np
>>> x = Tensor(np.arange(1, 4, dtype=np.int32))
>>> F.one_hot(x, num_classes=4)
Tensor([[0 1 0 0]
 [0 0 1 0]
 [0 0 0 1]], dtype=int32, device=xpux:0)
Return type

Tensor