import torch


def convert_to_one_hots(a, num_classes: int, dtype=torch.int, device=None):
    """
    Convert class index array (num_sample,) to an one hots array
    (num_sample, num_classes)

    Args:
        a: index array
        num_classes: number of classes
        dtype: data type

    Returns:
        one hots array in shape of (a.shape[0], num_classes)
    """
    one_hots = torch.zeros((len(a), num_classes), dtype=dtype, device=device)
    one_hots[torch.arange(one_hots.shape[0]), a] = 1
    return one_hots