Skip to content

GitLab

Commit 635b089a authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 335952206
parent fba8ea2f
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
orbit/utils/common.py
View file @ 635b089a
...@@ -63,7 +63,7 @@ def make_distributed_dataset(strategy, dataset_or_fn, *args, **kwargs): ...@@ -63,7 +63,7 @@ def make_distributed_dataset(strategy, dataset_or_fn, *args, **kwargs):
strategy = tf.distribute.get_strategy() strategy = tf.distribute.get_strategy()
if isinstance(dataset_or_fn, tf.data.Dataset): if isinstance(dataset_or_fn, tf.data.Dataset):
return strategy.distribute_dataset(dataset_or_fn) return strategy.experimental_distribute_dataset(dataset_or_fn)
if not callable(dataset_or_fn): if not callable(dataset_or_fn):
raise ValueError("`dataset_or_fn` should be either callable or an instance " raise ValueError("`dataset_or_fn` should be either callable or an instance "
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment