Skip to content

GitLab

Unverified Commit d9823dae authored by Haoyu Zhang's avatar Haoyu Zhang Committed by GitHub
Browse files

Add synthetic data monkey patch to OneDeviceStrategy as well (#6505)

parent cd63cf67
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
official/utils/misc/distribution_utils.py
View file @ d9823dae
...@@ -198,6 +198,7 @@ def _undo_monkey_patch_dataset_method(strategy): ...@@ -198,6 +198,7 @@ def _undo_monkey_patch_dataset_method(strategy):
def set_up_synthetic_data(): def set_up_synthetic_data():
_monkey_patch_dataset_method(tf.distribute.OneDeviceStrategy)
_monkey_patch_dataset_method(tf.distribute.MirroredStrategy) _monkey_patch_dataset_method(tf.distribute.MirroredStrategy)
# TODO(tobyboyd): Remove when contrib.distribute is all in core. # TODO(tobyboyd): Remove when contrib.distribute is all in core.
if hasattr(tf, 'contrib'): if hasattr(tf, 'contrib'):
...@@ -208,6 +209,7 @@ def set_up_synthetic_data(): ...@@ -208,6 +209,7 @@ def set_up_synthetic_data():
def undo_set_up_synthetic_data(): def undo_set_up_synthetic_data():
_undo_monkey_patch_dataset_method(tf.distribute.OneDeviceStrategy)
_undo_monkey_patch_dataset_method(tf.distribute.MirroredStrategy) _undo_monkey_patch_dataset_method(tf.distribute.MirroredStrategy)
# TODO(tobyboyd): Remove when contrib.distribute is all in core. # TODO(tobyboyd): Remove when contrib.distribute is all in core.
if hasattr(tf, 'contrib'): if hasattr(tf, 'contrib'):
......
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