From 0d91d6df6e01988fdf93769fde3f584aae9ab404 Mon Sep 17 00:00:00 2001 From: Bruno Arine Date: Thu, 15 Jun 2023 21:05:15 -0300 Subject: [PATCH] Add world_size check before setting up DDP train (#3191) --- ultralytics/yolo/engine/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ultralytics/yolo/engine/trainer.py b/ultralytics/yolo/engine/trainer.py index 5ef0d53..c3f8fcd 100644 --- a/ultralytics/yolo/engine/trainer.py +++ b/ultralytics/yolo/engine/trainer.py @@ -217,7 +217,7 @@ class BaseTrainer: callbacks_backup = callbacks.default_callbacks.copy() # backup callbacks as check_amp() resets them self.amp = torch.tensor(check_amp(self.model), device=self.device) callbacks.default_callbacks = callbacks_backup # restore callbacks - if RANK > -1: # DDP + if RANK > -1 and world_size > 1: # DDP dist.broadcast(self.amp, src=0) # broadcast the tensor from rank 0 to all other ranks (returns None) self.amp = bool(self.amp) # as boolean self.scaler = amp.GradScaler(enabled=self.amp)