mirror of
https://github.com/macaodha/batdetect2.git
synced 2026-01-10 00:59:34 +01:00
Small fix
This commit is contained in:
parent
9d1497b3f4
commit
98f83e8b34
@ -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,
|
||||
|
||||
Loading…
Reference in New Issue
Block a user