def get_num_classes(DATASET: str) -> int:
    """Get the number of classes to be classified by dataset."""
    if DATASET == 'MELD':
        NUM_CLASSES = 7
    elif DATASET == 'IEMOCAP':
        NUM_CLASSES = 6
    else:
        raise ValueError

    return NUM_CLASSES