import torch
from typing import Optional


def torch_one_hot(target: torch.Tensor, num_classes: Optional[int] = None) -> torch.Tensor:
    """
    Compute one hot encoding of input tensor

    Args:
        target: tensor to be converted
        num_classes: number of classes. If :attr:`num_classes` is None,
            the maximum of target is used

    Returns:
        torch.Tensor: one hot encoded tensor
    """
    if num_classes is None:
        num_classes = int(target.max().detach().item() + 1)
    dtype, device = target.dtype, target.device
    target_onehot = torch.zeros(*target.shape, num_classes, dtype=dtype, device=device)
    return target_onehot.scatter_(1, target.unsqueeze_(1), 1.0)