fix ignore to dice (#588)

Deci-AI · Dec 28, 2022 · f460027 · f460027
1 parent 24772f4
commit f460027
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/super_gradients/training/losses/dice_ce_edge_loss.py b/src/super_gradients/training/losses/dice_ce_edge_loss.py
@@ -59,7 +59,7 @@ def __init__(
             self.binary_dice = BinaryDiceLoss(apply_sigmoid=True)
 
         self.ce_edge = MaskAttentionLoss(criterion=nn.CrossEntropyLoss(reduction="none", ignore_index=ignore_index), loss_weights=ce_edge_weights)
-        self.dice_loss = DiceLoss(apply_softmax=True, ignore_index=ignore_index)
+        self.dice_loss = DiceLoss(apply_softmax=True, ignore_index=None if ignore_index < 0 else ignore_index)
 
     @property
     def component_names(self):