import torch


def to_one_hot(tensor: torch.Tensor, n_dims=None) -> torch.Tensor:
    """
    Take long `tensor` with `n_dims` and
    convert it to one-hot representation with n + 1 dimensions.
    """
    tensor = tensor.long().view(-1, 1)
    n_dims = n_dims if n_dims else int(torch.max(tensor)) + 1
    one_hot_tensor = torch.zeros(tensor.size(0), n_dims).scatter_(1, tensor, 1)
    return one_hot_tensor.view(tensor.size(0), -1)