fix: move kpts mAP output to correct device

luxonis · Dec 16, 2024 · 0c45b8c · 0c45b8c
1 parent 6445cbf
commit 0c45b8c
Showing 1 changed file with 29 additions and 10 deletions.
diff --git a/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py b/luxonis_train/attached_modules/metrics/mean_average_precision_keypoints.py
@@ -254,17 +254,36 @@ def compute(self) -> tuple[Tensor, dict[str, Tensor]]:
             self.coco_eval.summarize()
             stats = self.coco_eval.stats
 
-        kpt_map = torch.tensor([stats[0]], dtype=torch.float32)
+        device = self.pred_keypoints[0].device
+        kpt_map = torch.tensor([stats[0]], dtype=torch.float32, device=device)
         return kpt_map, {
-            "kpt_map_50": torch.tensor([stats[1]], dtype=torch.float32),
-            "kpt_map_75": torch.tensor([stats[2]], dtype=torch.float32),
-            "kpt_map_medium": torch.tensor([stats[3]], dtype=torch.float32),
-            "kpt_map_large": torch.tensor([stats[4]], dtype=torch.float32),
-            "kpt_mar": torch.tensor([stats[5]], dtype=torch.float32),
-            "kpt_mar_50": torch.tensor([stats[6]], dtype=torch.float32),
-            "kpt_mar_75": torch.tensor([stats[7]], dtype=torch.float32),
-            "kpt_mar_medium": torch.tensor([stats[8]], dtype=torch.float32),
-            "kpt_mar_large": torch.tensor([stats[9]], dtype=torch.float32),
+            "kpt_map_50": torch.tensor(
+                [stats[1]], dtype=torch.float32, device=device
+            ),
+            "kpt_map_75": torch.tensor(
+                [stats[2]], dtype=torch.float32, device=device
+            ),
+            "kpt_map_medium": torch.tensor(
+                [stats[3]], dtype=torch.float32, device=device
+            ),
+            "kpt_map_large": torch.tensor(
+                [stats[4]], dtype=torch.float32, device=device
+            ),
+            "kpt_mar": torch.tensor(
+                [stats[5]], dtype=torch.float32, device=device
+            ),
+            "kpt_mar_50": torch.tensor(
+                [stats[6]], dtype=torch.float32, device=device
+            ),
+            "kpt_mar_75": torch.tensor(
+                [stats[7]], dtype=torch.float32, device=device
+            ),
+            "kpt_mar_medium": torch.tensor(
+                [stats[8]], dtype=torch.float32, device=device
+            ),
+            "kpt_mar_large": torch.tensor(
+                [stats[9]], dtype=torch.float32, device=device
+            ),
         }
 
     def _get_coco_format(