From ea699a53ea95d90420427a95e9d5ee271940fac1 Mon Sep 17 00:00:00 2001 From: Bhaskar Gautam Date: Tue, 18 Aug 2020 09:33:23 +0530 Subject: [PATCH] Added support for TF 2.3.0 --- model_training_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/model_training_utils.py b/model_training_utils.py index f947fd7..c118311 100644 --- a/model_training_utils.py +++ b/model_training_utils.py @@ -236,7 +236,7 @@ def train_steps(iterator, steps): 'retracing.') for _ in tf.range(steps): - strategy.experimental_run_v2(_replicated_step, args=(next(iterator),)) + strategy.run(_replicated_step, args=(next(iterator),)) def train_single_step(iterator): """Performs a distributed training step. @@ -247,7 +247,7 @@ def train_single_step(iterator): Raises: ValueError: Any of the arguments or tensor shapes are invalid. """ - strategy.experimental_run_v2(_replicated_step, args=(next(iterator),)) + strategy.run(_replicated_step, args=(next(iterator),)) def test_step(iterator): """Calculates evaluation metrics on distributed devices.""" @@ -260,7 +260,7 @@ def _test_step_fn(inputs): for metric in eval_metrics: metric.update_state(labels, model_outputs) - strategy.experimental_run_v2(_test_step_fn, args=(next(iterator),)) + strategy.run(_replicated_step, args=(next(iterator),)) if not run_eagerly: train_single_step = tf.function(train_single_step)