Skip to content

GitLab

Unverified Commit 325dd761 authored by Haoyu Zhang's avatar Haoyu Zhang Committed by GitHub
Browse files

Specify NCCL as the all reduce algorithm (#6662)

parent 37f17256
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
official/resnet/keras/keras_imagenet_main.py
View file @ 325dd761
...@@ -124,7 +124,8 @@ def run(flags_obj): ...@@ -124,7 +124,8 @@ def run(flags_obj):
strategy = distribution_utils.get_distribution_strategy( strategy = distribution_utils.get_distribution_strategy(
distribution_strategy=flags_obj.distribution_strategy, distribution_strategy=flags_obj.distribution_strategy,
num_gpus=flags_obj.num_gpus, num_gpus=flags_obj.num_gpus,
num_workers=distribution_utils.configure_cluster()) num_workers=distribution_utils.configure_cluster(),
all_reduce_alg='nccl')
strategy_scope = distribution_utils.get_strategy_scope(strategy) strategy_scope = distribution_utils.get_strategy_scope(strategy)
......
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