Fix central storage 2 gpus combination

PiperOrigin-RevId: 336741815
Change-Id: I047c48c9d3fd34d25ef97994e244b15e8dfda738
This commit is contained in:
Ran Chen 2020-10-12 14:18:49 -07:00 committed by TensorFlower Gardener
parent 30429491d7
commit e5a3292a48

View File

@ -218,7 +218,7 @@ mirrored_strategy_with_cpu_1_and_2.__doc__ = (
""")
central_storage_strategy_with_two_gpus = combinations.NamedDistribution(
"CentralStorage2GPUs",
lambda: CentralStorageStrategy._from_num_gpus(2), # pylint: disable=protected-access
lambda: CentralStorageStrategy(["/gpu:0", "/gpu:1"]),
required_gpus=2)
central_storage_strategy_with_gpu_and_cpu = combinations.NamedDistribution(
"CentralStorageCPUAndGPU",