Validate calibration loss weights

Signed-off-by: Theo Barfoot <theo.barfoot@gmail.com>

Author: theo-barfoot

monai/losses/calibration.py

@@ -109,6 +109,8 @@ def __init__(
         if class_weight is not None:
             if class_weight.ndim > 1:
                 raise ValueError("weight must be a scalar or a one-dimensional sequence.")
+            if not torch.all(torch.isfinite(class_weight)):
+                raise ValueError("weight must contain only finite values.")
             if torch.any(class_weight < 0):
                 raise ValueError("the value/values of the `weight` should be no less than 0.")
 
@@ -353,8 +355,8 @@ def __init__(
         right: bool = False,
         ignore_empty_classes: bool = True,
     ) -> None:
-        if empty_weight < 0:
-            raise ValueError(f"empty_weight must be >= 0, got {empty_weight}.")
+        if not 0 <= empty_weight < float("inf"):
+            raise ValueError(f"empty_weight must be finite and >= 0, got {empty_weight}.")
         super().__init__(
             num_bins,
             include_background,

tests/losses/test_calibration_loss.py

@@ -159,6 +159,14 @@ def test_invalid_options_and_shapes(self):
                     loss_type(other_act=1)  # type: ignore[arg-type]
                 with self.assertRaises(ValueError):
                     loss_type(weight=[1.0, -1.0])
+                for invalid_weight in (float("nan"), float("inf"), -float("inf")):
+                    with self.subTest(invalid_weight=invalid_weight):
+                        with self.assertRaises(ValueError):
+                            loss_type(weight=invalid_weight)
+                        with self.assertRaises(ValueError):
+                            loss_type(weight=[1.0, invalid_weight])
+                        with self.assertRaises(ValueError):
+                            loss_type(weight=torch.tensor([invalid_weight]))
                 with self.assertRaises(ValueError):
                     loss_type(weight=[[1.0]])
                 with self.assertRaises(ValueError):
@@ -171,8 +179,9 @@ def test_invalid_options_and_shapes(self):
                     loss_type()(torch.ones(1, 2), torch.ones(1, 2))
                 with self.assertRaises(TypeError):
                     loss_type()(torch.ones(1, 2, 2, dtype=torch.int64), torch.ones(1, 2, 2))
-        with self.assertRaises(ValueError):
-            SoftL1ACELoss(empty_weight=-1)
+        for invalid_empty_weight in (-1.0, float("nan"), float("inf"), -float("inf")):
+            with self.subTest(invalid_empty_weight=invalid_empty_weight), self.assertRaises(ValueError):
+                SoftL1ACELoss(empty_weight=invalid_empty_weight)
 
     def test_single_channel_warnings(self):
         prediction = torch.tensor([[[0.2, 0.8]]])