Add criterion for checking valid weighing factor

computational-cell-analytics · Jul 10, 2024 · 47b191f · 47b191f
1 parent f2bd62e
commit 47b191f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/micro_sam/training/semantic_sam_trainer.py b/micro_sam/training/semantic_sam_trainer.py
@@ -51,6 +51,10 @@ def __init__(
         self.num_classes = num_classes
         self.compute_ce_loss = nn.CrossEntropyLoss()
         self.dice_weight = dice_weight
+
+        if self.dice_weight is not None:
+            assert self.dice_weight > 0 and self.dice_weight < 1, "The weight factor should lie between 0 and 1."
+
         self._kwargs = kwargs
 
     def _compute_loss(self, y, masks):