From 470b120a1bac290bb1b98b1c9ebad3bd05e2dcfc Mon Sep 17 00:00:00 2001
From: Chia-Hsiang Tsai <84863554+Tsai-chia-hsiang@users.noreply.github.com>
Date: Thu, 11 Jul 2024 01:59:16 +0800
Subject: [PATCH] BaseTrainer with `find_unused_parameters=True` when using
 DistributedDataParallel() (#14323)

---
 ultralytics/engine/trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index c833e76165..3fb3e0b852 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)