import torch


def one_hot(y, num_dim=10):
    """
    One Hot Encoding, similar to `torch.eye(num_dim).index_select(dim=0, index=y)`
    :param y: N-dim tenser
    :param num_dim: do one-hot labeling from `0` to `num_dim-1`
    :return: shape = (batch_size, num_dim)
    """
    one_hot_y = torch.zeros(y.size(0), num_dim)
    if y.is_cuda:
        one_hot_y = one_hot_y.cuda()
    return one_hot_y.scatter_(1, y.view(-1, 1), 1.)