Skip to content

GitLab

Unverified Commit 98672351 authored by Yuefeng Zhou's avatar Yuefeng Zhou Committed by GitHub
Browse files

Update distribution_utils.py (#6615)

parent 4ad73a1c
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 3 deletions
official/utils/misc/distribution_utils.py
View file @ 98672351
...@@ -230,9 +230,8 @@ def configure_cluster(worker_hosts=None, task_index=-1): ...@@ -230,9 +230,8 @@ def configure_cluster(worker_hosts=None, task_index=-1):
""" """
tf_config = json.loads(os.environ.get('TF_CONFIG', '{}')) tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
if tf_config: if tf_config:
num_workers = len(tf_config['cluster']['worker']) num_workers = (len(tf_config['cluster'].get('chief', [])) +
if tf_config['cluster'].get('chief', None): len(tf_config['cluster'].get('worker', [])))
num_workers += 1
elif worker_hosts: elif worker_hosts:
workers = worker_hosts.split(',') workers = worker_hosts.split(',')
num_workers = len(workers) num_workers = len(workers)
......
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