From 98f83e8b347a2d82e2dd67d336315b1631b5b3fa Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Mon, 18 Aug 2025 10:25:14 +0100 Subject: [PATCH] Small fix --- src/batdetect2/models/detectors.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/batdetect2/models/detectors.py b/src/batdetect2/models/detectors.py index c5ab691..854f71e 100644 --- a/src/batdetect2/models/detectors.py +++ b/src/batdetect2/models/detectors.py @@ -114,8 +114,8 @@ class Detector(DetectionModel): `(B, C_out, H, W)`. """ features = self.backbone(spec) - detection = self.detector_head(features) classification = self.classifier_head(features) + detection = classification.sum(dim=1, keep_dim=True) size_preds = self.bbox_head(features) return ModelOutput( detection_probs=detection,