def get_dataset_params(dataset):
    """
    Get the capacities to test for the particular dataset.
    :param dataset: string of dataset
    :return: list of capacities
    """
    if dataset == 'cub200':
        capacity_list = [2, 4, 8, 16]
    else:
        raise NotImplementedError('Dataset not supported.')
    return capacity_list