From 2296f1546fe252d7293b48ffb8e192d1e5f2a85b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Thu, 17 Jun 2021 23:24:30 +0200
Subject: [PATCH] Update `WORLD_SIZE` and `RANK` retrieval (#3670)

---
 train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/train.py b/train.py
index 9d71e7056800..3eb866345d47 100644
--- a/train.py
+++ b/train.py
@@ -502,8 +502,8 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
     opt = parser.parse_args()
 
     # Set DDP variables
-    opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
-    opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
+    opt.world_size = int(getattr(os.environ, 'WORLD_SIZE', 1))
+    opt.global_rank = int(getattr(os.environ, 'RANK', -1))
     set_logging(opt.global_rank)
     if opt.global_rank in [-1, 0]:
         check_git_status()