diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py index c833e761..3fb3e0b8 100644 --- a/ultralytics/engine/trainer.py +++ b/ultralytics/engine/trainer.py @@ -266,7 +266,7 @@ class BaseTrainer: self.amp = bool(self.amp) # as boolean self.scaler = torch.cuda.amp.GradScaler(enabled=self.amp) if world_size > 1: - self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[RANK]) + self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[RANK], find_unused_parameters=True) # Check imgsz gs = max(int(self.model.stride.max() if hasattr(self.model, "stride") else 32), 32) # grid size (max stride)