Update WORLD_SIZE and RANK retrieval (#3670)

ultralytics · Jun 17, 2021 · 2296f15 · 2296f15
1 parent 045d5d8
commit 2296f15
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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()