Fixed Retrieving Labels in AdaptiveDetectionLoss (#59)

luxonis · Oct 9, 2024 · 6f5a9c3 · 6f5a9c3
1 parent 575e0ca
commit 6f5a9c3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/luxonis_train/attached_modules/losses/adaptive_detection_loss.py b/luxonis_train/attached_modules/losses/adaptive_detection_loss.py
@@ -94,7 +94,7 @@ def prepare(
         batch_size = pred_scores.shape[0]
         device = pred_scores.device
 
-        target = labels[self.task][0].to(device)
+        target = self.get_label(labels)[0]
         if self.gt_bboxes_scale is None:
             self.gt_bboxes_scale = torch.tensor(
                 [