cleanup destroy

ultralytics · Jun 18, 2021 · 9c4ac05 · 9c4ac05
1 parent b09b415
commit 9c4ac05
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/train.py b/train.py
@@ -459,7 +459,8 @@ def train(hyp,  # path/to/hyp.yaml or hyp dictionary
                                                 name='run_' + wandb_logger.wandb_run.id + '_model',
                                                 aliases=['latest', 'best', 'stripped'])
         wandb_logger.finish_run()
-    else:
+
+    if WORLD_SIZE > 1:
         dist.destroy_process_group()
     torch.cuda.empty_cache()
     return results