fix: ddp bug

luxonis · Dec 13, 2024 · 6445cbf · 6445cbf
1 parent ec717aa
commit 6445cbf
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision.py b/luxonis_train/attached_modules/metrics/mean_average_precision.py
@@ -83,5 +83,7 @@ def compute(self) -> tuple[Tensor, dict[str, Tensor]]:
                     )
 
         map = metric_dict.pop("map")
-
+        # WARNING: fix DDP pl.log error
+        map = map.to(self.device)
+        metric_dict = {k: v.to(self.device) for k, v in metric_dict.items()}
         return map, metric_dict