import torch
from typing import Optional


def one_hot(target: torch.Tensor, num_classes: Optional[int] = None) -> torch.Tensor:
    """source: https://github.com/PhoenixDL/rising. Computes one-hot encoding of input tensor.
    Args:
        target (torch.Tensor): tensor to be converted
        num_classes (Optional[int], optional): number of classes. If None, the maximum value of target is used.
    Returns:
        torch.Tensor: one-hot encoded tensor of the target
    """
    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)