Function iterator_function_def

dali/test/python/jax_plugin/utils.py:116–125

(shard_id=0, num_shards=1)

def iterator_function_def(shard_id=0, num_shards=1):
    files, labels = fn.readers.file(
        name="reader",
        files=file_names,
        labels=file_labels,
        shard_id=shard_id,
        num_shards=num_shards,
    )

    return labels.gpu()