def get_num_classes(dataset_name: str) -> int:
    """
    Get the number of supervised loss given a dataset name.
    :param dataset_name: dataset name
    :return: number of supervised class.
    """
    if dataset_name == "cifar10":
        return 10
    elif "cifar100" in dataset_name:
        return 100
    else:
        raise ValueError("Supported datasets are only original CIFAR10/100")