Skip to content

GitLab

Unverified Commit b66ef95e authored by Toby Boyd's avatar Toby Boyd Committed by GitHub
Browse files

One device strat (#6196) 

* One device from contrib to core.

* remove test code.
parent 589ac399
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
official/utils/misc/distribution_utils.py
View file @ b66ef95e
...@@ -59,12 +59,12 @@ def get_distribution_strategy(distribution_strategy="default", ...@@ -59,12 +59,12 @@ def get_distribution_strategy(distribution_strategy="default",
if (distribution_strategy == "one_device" or if (distribution_strategy == "one_device" or
(distribution_strategy == "default" and num_gpus <= 1)): (distribution_strategy == "default" and num_gpus <= 1)):
if num_gpus == 0: if num_gpus == 0:
return tf.contrib.distribute.OneDeviceStrategy("device:CPU:0") return tf.distribute.OneDeviceStrategy("device:CPU:0")
else: else:
if num_gpus > 1: if num_gpus > 1:
raise ValueError("`OneDeviceStrategy` can not be used for more than " raise ValueError("`OneDeviceStrategy` can not be used for more than "
"one device.") "one device.")
return tf.contrib.distribute.OneDeviceStrategy("device:GPU:0") return tf.distribute.OneDeviceStrategy("device:GPU:0")
if distribution_strategy in ("mirrored", "default"): if distribution_strategy in ("mirrored", "default"):
if num_gpus == 0: if num_gpus == 0:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment