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:
parent
d6a7f2d044
commit
e4ec26f50c
@ -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
|
||||
|
@ -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(
|
||||
|
@ -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"])
|
||||
|
||||
|
||||
|
@ -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
|
||||
]
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user