Skip to content

GitLab

"...orqa/git@developer.sourcefind.cn:OpenDAS/megatron-lm.git" did not exist on "0024a5c66f90c7d3d02f7ef08a773aace6deb155"
Commit 122bb012 authored by Toby Boyd's avatar Toby Boyd
Browse files

fix lint errors.

parent d945e99f
Show whitespace changes
Inline Side-by-side
Showing with 11 additions and 9 deletions
official/utils/misc/distribution_utils.py
View file @ 122bb012
...@@ -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:
...@@ -52,8 +53,9 @@ def get_distribution_strategy( ...@@ -52,8 +53,9 @@ def get_distribution_strategy(
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"
"True.".format(num_gpus))
else: # num_gpus > 1 and not turn_off_distribution_strategy 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(
......
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