diff --git a/tensorflow/python/keras/distribute/distribute_strategy_test.py b/tensorflow/python/keras/distribute/distribute_strategy_test.py index bbd62c11cb3..6c3099ed524 100644 --- a/tensorflow/python/keras/distribute/distribute_strategy_test.py +++ b/tensorflow/python/keras/distribute/distribute_strategy_test.py @@ -237,8 +237,7 @@ strategies_minus_tpu = [ ] tpu_strategies = [ - strategy_combinations.tpu_strategy, # steps_per_run=2 - strategy_combinations.tpu_strategy_one_step + strategy_combinations.tpu_strategy, ] all_strategies = strategies_minus_tpu + tpu_strategies diff --git a/tensorflow/python/keras/distribute/keras_correctness_test_base.py b/tensorflow/python/keras/distribute/keras_correctness_test_base.py index 825e94b9eba..8f019cdb7c9 100644 --- a/tensorflow/python/keras/distribute/keras_correctness_test_base.py +++ b/tensorflow/python/keras/distribute/keras_correctness_test_base.py @@ -51,7 +51,6 @@ all_strategies = [ strategy_combinations.mirrored_strategy_with_gpu_and_cpu, strategy_combinations.mirrored_strategy_with_two_gpus, strategy_combinations.tpu_strategy, # steps_per_run=2 - strategy_combinations.tpu_strategy_one_step, ] @@ -118,7 +117,6 @@ def test_combinations_for_embedding_model(): def test_combinations_with_tpu_strategies(): tpu_strategies = [ strategy_combinations.tpu_strategy, - strategy_combinations.tpu_strategy_one_step ] return (combinations.times( diff --git a/tensorflow/python/keras/distribute/keras_metrics_test.py b/tensorflow/python/keras/distribute/keras_metrics_test.py index af82b79066c..bf592898d00 100644 --- a/tensorflow/python/keras/distribute/keras_metrics_test.py +++ b/tensorflow/python/keras/distribute/keras_metrics_test.py @@ -84,9 +84,7 @@ def all_combinations(): def tpu_combinations(): return combinations.combine( - distribution=[ - strategy_combinations.tpu_strategy_one_step, - ], + distribution=[strategy_combinations.tpu_strategy,], mode=["graph"]) diff --git a/tensorflow/python/keras/distribute/keras_stateful_lstm_model_correctness_test.py b/tensorflow/python/keras/distribute/keras_stateful_lstm_model_correctness_test.py index f6344bd3646..199a1f390a4 100644 --- a/tensorflow/python/keras/distribute/keras_stateful_lstm_model_correctness_test.py +++ b/tensorflow/python/keras/distribute/keras_stateful_lstm_model_correctness_test.py @@ -32,7 +32,6 @@ def strategies_for_stateful_embedding_model(): return [ strategy_combinations.tpu_strategy_one_core, - strategy_combinations.tpu_strategy_one_step_one_core ]