From 4a8d567eb494022443ad28f1e52ae5a908d649c0 Mon Sep 17 00:00:00 2001 From: Alessio Placitelli Date: Fri, 10 Oct 2025 09:59:45 +0200 Subject: [PATCH] Enforce long type for labels when using one class Without this change, when training a model with a single class, the training will fail as `torch.zeros` would fill `labels` with floating point `0.0`, which can't be used as indices. --- opentad/models/detectors/single_stage.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/opentad/models/detectors/single_stage.py b/opentad/models/detectors/single_stage.py index 3ba72941..4e871193 100644 --- a/opentad/models/detectors/single_stage.py +++ b/opentad/models/detectors/single_stage.py @@ -109,7 +109,7 @@ def post_processing(self, predictions, metas, post_cfg, ext_cls, **kwargs): if num_classes == 1: scores = scores.squeeze(-1) - labels = torch.zeros(scores.shape[0]).contiguous() + labels = torch.zeros(scores.shape[0], dtype=torch.long).contiguous() else: pred_prob = scores.flatten() # [N*class]