def METHOD_NAME(
    name,
    batch_size,
    num_epochs,
    split,
    framework,
    dataset_dir,
    shuffle_files=False,
    preprocessing_fn=None,
):
    # Instance types based on framework
    instance_types = {
        "pytorch": torch.utils.data.DataLoader,
        "tf": (tf.compat.v2.data.Dataset, tf.compat.v1.data.Dataset),
        "numpy": ArmoryDataGenerator,
    }
    if framework not in instance_types:
        raise Exception(
            "Unrecognized Armory Dataset Framework: {}".format(framework)
        )
    ds = getattr(datasets, name)
    dataset = ds(
        split=split,
        batch_size=batch_size,
        epochs=num_epochs,
        dataset_dir=dataset_dir,
        framework=framework,
        shuffle_files=shuffle_files,
        preprocessing_fn=preprocessing_fn,
    )
    print(type(dataset))
    assert isinstance(dataset, instance_types[framework])
    if framework == "numpy":
        assert dataset.batch_size == batch_size
        assert dataset.batches_per_epoch == (
            dataset.size // batch_size + bool(dataset.size % batch_size)
        )
    return dataset