From 6445cbff6b93c5ca0d877371ce7031000ab731da Mon Sep 17 00:00:00 2001 From: Jernej Sabadin Date: Fri, 13 Dec 2024 11:57:39 +0100 Subject: [PATCH] fix: ddp bug --- .../attached_modules/metrics/mean_average_precision.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision.py b/luxonis_train/attached_modules/metrics/mean_average_precision.py index 56937115..d4731988 100644 --- 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