import torch


def one_hot(index, size):
    """
    Expand indexes into one-hot vectors.

    Parameters:
        index (Tensor): index
        size (int): size of the one-hot dimension
    """
    shape = list(index.shape) + [size]
    result = torch.zeros(shape, device=index.device)
    if index.numel():
        assert index.min() >= 0
        assert index.max() < size
        result.scatter_(-1, index.unsqueeze(-1), 1)
    return result