import torch


def one_hot(index, n_cat, dropColumn=False):
    """
    expects tensor of shape (n_samples,1), returns one-hot array size
    (n_samples, n_cat).

    Optionally, can drop the first column of the encoding to prevent colinearity
    among the predictors. This is only necessary if you care about the values
    of the inferred parameters during inference.
    """

    onehot = torch.zeros(index.size(0), n_cat, device=index.device)
    onehot.scatter_(1, index.type(torch.long), 1)
    if dropColumn:
        return onehot.type(torch.float32)[:,1:]
    else:
        return onehot.type(torch.float32)