import torch


def one_hot_encoder(idx, n_cls):
    """
    One hot encoder for categorical features

    Args:
        idx: index of the categorical feature
        n_cls: number of classes

    Returns:
        one hot encoded tensor
    """
    if idx.dim() == 1:
        idx = idx.unsqueeze(1)
    onehot = torch.zeros(idx.size(0), n_cls)
    onehot = onehot.to(idx.device)
    onehot.scatter_(1, idx.long(), 1)
    return onehot