Update distribution_utils.py (#6615)

98672351 · Yuefeng Zhou · GitHub · 4ad73a1c · 98672351
Unverified Commit 98672351 authored Apr 24, 2019 by Yuefeng Zhou Committed by GitHub Apr 24, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 3 deletions

official/utils/misc/distribution_utils.py official/utils/misc/distribution_utils.py +2 -3

No files found.
--- a/official/utils/misc/distribution_utils.py
+++ b/official/utils/misc/distribution_utils.py
@@ -230,9 +230,8 @@ def configure_cluster(worker_hosts=None, task_index=-1):
  """
  tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
  if tf_config:
-    num_workers = len(tf_config['cluster']['worker'])
-    if tf_config['cluster'].get('chief', None):
-      num_workers += 1
+    num_workers = (len(tf_config['cluster'].get('chief', [])) +
+                   len(tf_config['cluster'].get('worker', [])))
  elif worker_hosts:
    workers = worker_hosts.split(',')
    num_workers = len(workers)