From 8f17df9ae9f15b7347128ea5a098dab268bb3b01 Mon Sep 17 00:00:00 2001 From: Yuexin Wu Date: Wed, 17 May 2023 11:33:47 -0700 Subject: [PATCH] Disable preemption_watcher as it's not supported externally. PiperOrigin-RevId: 532853592 --- official/nlp/train.py | 4 +++- official/vision/train.py | 4 +++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/official/nlp/train.py b/official/nlp/train.py index 621d095c2..3bdcaf3f5 100644 --- a/official/nlp/train.py +++ b/official/nlp/train.py @@ -53,7 +53,9 @@ def _run_experiment_with_preemption_recovery(params, model_dir): **params.runtime.model_parallelism()) with distribution_strategy.scope(): task = task_factory.get_task(params.task, logging_dir=model_dir) - preemption_watcher = tf.distribute.experimental.PreemptionWatcher() + # pylint: disable=line-too-long + preemption_watcher = None # copybara-replace + # pylint: enable=line-too-long train_lib.run_experiment( distribution_strategy=distribution_strategy, diff --git a/official/vision/train.py b/official/vision/train.py index d391d7aea..2b87a18f4 100644 --- a/official/vision/train.py +++ b/official/vision/train.py @@ -46,7 +46,9 @@ def _run_experiment_with_preemption_recovery(params, model_dir): tpu_address=params.runtime.tpu) with distribution_strategy.scope(): task = task_factory.get_task(params.task, logging_dir=model_dir) - preemption_watcher = tf.distribute.experimental.PreemptionWatcher() + # pylint: disable=line-too-long + preemption_watcher = None # copybara-replace + # pylint: enable=line-too-long train_lib.run_experiment( distribution_strategy=distribution_strategy, -- GitLab