def get_input_size(dataset_name):
    """
    Returns the size of input
    """
    dataset_name = dataset_name.lower()
    mapping = {
        'mnist': 28 * 28,
        'adult': 6,
        'pums': 4,
        'power': 8 
    }
    if dataset_name not in mapping:
        err_msg = f"Unknown dataset '{dataset_name}'. Please choose one in {list(mapping.keys())}."
        raise ValueError(err_msg)

    return mapping[dataset_name]