Update keras/distribute test wrt to tpu_strategy_one_step{_one_core}

The step argument has been removed in TF2.

PiperOrigin-RevId: 332477799
Change-Id: Ifcfa8bba7228baad73ba3993a9024031bee70fe0
This commit is contained in:
Scott Zhu 2020-09-18 10:29:08 -07:00 committed by TensorFlower Gardener
parent d6a7f2d044
commit e4ec26f50c
4 changed files with 2 additions and 8 deletions

View File

@ -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

View File

@ -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(

View File

@ -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"])

View File

@ -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
]