orbit.utils.make_distributed_dataset

A utility function to help create a tf.distribute.DistributedDataset.

orbit.utils.make_distributed_dataset(
    strategy, dataset_or_fn, *args, **kwargs
)

strategy An instance of tf.distribute.Strategy.
dataset_or_fn A instance of tf.data.Dataset, or a "dataset function" returning a tf.data.Dataset. If it is a function, it may optionally have an argument named input_context which will be passed a tf.distribute.InputContext instance.
*args Any positional arguments to pass through to dataset_or_fn.
**kwargs Any keyword arguments to pass through to dataset_or_fn, except that the input_options keyword is used to specify a tf.distribute.InputOptions for making the distributed dataset.

A distributed Dataset.