fix: smooth bce loss generics

luxonis · Jan 29, 2025 · 32f4bfc · 32f4bfc
1 parent 44f40ba
commit 32f4bfc
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/luxonis_train/attached_modules/losses/smooth_bce_with_logits.py b/luxonis_train/attached_modules/losses/smooth_bce_with_logits.py
@@ -9,7 +9,7 @@
 from .bce_with_logits import BCEWithLogitsLoss
 
 
-class SmoothBCEWithLogitsLoss(BaseLoss[list[Tensor], Tensor]):
+class SmoothBCEWithLogitsLoss(BaseLoss[Tensor, Tensor]):
     supported_tasks: list[TaskType] = [
         TaskType.SEGMENTATION,
         TaskType.CLASSIFICATION,
@@ -65,6 +65,7 @@ def forward(self, predictions: Tensor, target: Tensor) -> Tensor:
         @rtype: Tensor
         @return: A scalar tensor.
         """
+        print(predictions[0].shape)
         if predictions.shape != target.shape:
             raise RuntimeError(
                 f"Target tensor dimension ({target.shape}) and predictions tensor "