fix export bug

luxonis · Jan 16, 2025 · 04a6ab6 · 04a6ab6
1 parent bcd40b8
commit 04a6ab6
Showing 1 changed file with 6 additions and 12 deletions.
diff --git a/luxonis_train/nodes/heads/fomo_head.py b/luxonis_train/nodes/heads/fomo_head.py
@@ -89,18 +89,12 @@ def _apply_nms_if_needed(self, heatmap: Tensor) -> Tensor:
         if not self.use_nms:
             return heatmap
 
-        return (
-            F.max_pool2d(
-                heatmap.unsqueeze(0).unsqueeze(
-                    0
-                ),  # Add dummy batch/channel dimensions
-                kernel_size=3,
-                stride=1,
-                padding=1,
-            )
-            .squeeze(0)
-            .squeeze(0)
-        )  # Remove dummy dimensions
+        return F.max_pool2d(
+            heatmap,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
 
     def _heatmap_to_kpts(self, heatmap: Tensor) -> List[Tensor]:
         """Convert heatmap to keypoint pairs using local-max NMS so that