import torch


def one_hot(class_vec, num_classes, dtype=None, device=None):
    """Create one-hot encoded tensor.

    ----
    Paramaters:
        class_vec: (N x 1) LongTensor
        num_classes: C
        dtype: data type for Tensor
        device: device for Tensor
    """
    I = torch.eye(num_classes, dtype=dtype, device=device)
    return I[class_vec.squeeze()]