Commit 122bb012 authored by Toby Boyd's avatar Toby Boyd
Browse files

fix lint errors.

parent d945e99f
...@@ -21,8 +21,9 @@ from __future__ import print_function ...@@ -21,8 +21,9 @@ from __future__ import print_function
import tensorflow as tf import tensorflow as tf
def get_distribution_strategy( def get_distribution_strategy(num_gpus,
num_gpus, all_reduce_alg=None, turn_off_distribution_strategy=False): all_reduce_alg=None,
turn_off_distribution_strategy=False):
"""Return a DistributionStrategy for running the model. """Return a DistributionStrategy for running the model.
Args: Args:
...@@ -42,19 +43,20 @@ def get_distribution_strategy( ...@@ -42,19 +43,20 @@ def get_distribution_strategy(
larger than 1 larger than 1
""" """
if num_gpus == 0: if num_gpus == 0:
if turn_off_distribution_strategy: if turn_off_distribution_strategy:
return None return None
else: else:
return tf.contrib.distribute.OneDeviceStrategy("device:CPU:0") return tf.contrib.distribute.OneDeviceStrategy("device:CPU:0")
elif num_gpus == 1: elif num_gpus == 1:
if turn_off_distribution_strategy: if turn_off_distribution_strategy:
return None return None
else: else:
return tf.contrib.distribute.OneDeviceStrategy("device:GPU:0") return tf.contrib.distribute.OneDeviceStrategy("device:GPU:0")
elif turn_off_distribution_strategy: elif turn_off_distribution_strategy:
raise ValueError("When %d GPUs are specified, turn_off_distribution_strategy" raise ValueError("When {} GPUs are specified, "
" flag cannot be set to True.".format(num_gpus)) "turn_off_distribution_strategy flag cannot be set to"
else: # num_gpus > 1 and not turn_off_distribution_strategy "True.".format(num_gpus))
else: # num_gpus > 1 and not turn_off_distribution_strategy
if all_reduce_alg: if all_reduce_alg:
return tf.contrib.distribute.MirroredStrategy( return tf.contrib.distribute.MirroredStrategy(
num_gpus=num_gpus, num_gpus=num_gpus,
......
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