Fix torch.export compatibility in loss functions

Replace torch.arange().tolist() with list(range()) in dice, focal,
and tversky losses to avoid graph breaks under torch.export tracing.
Move class_weight validation to __init__ in hausdorff_loss.

Signed-off-by: Soumya Snigdha Kundu <soumya_snigdha.kundu@kcl.ac.uk>

Author: aymuos15

monai/losses/dice.py

@@ -123,6 +123,8 @@ def __init__(
         self.smooth_dr = float(smooth_dr)
         self.batch = batch
         weight = torch.as_tensor(weight) if weight is not None else None
+        if weight is not None and weight.min() < 0:
+            raise ValueError("the value/values of the `weight` should be no less than 0.")
         self.register_buffer("class_weight", weight)
         self.class_weight: None | torch.Tensor
         self.soft_label = soft_label
@@ -181,7 +183,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             raise AssertionError(f"ground truth has different shape ({target.shape}) from input ({input.shape})")
 
         # reducing only spatial dimensions (not batch nor channels)
-        reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
+        reduce_axis: list[int] = list(range(2, len(input.shape)))
         if self.batch:
             # reducing spatial dimensions and batch
             reduce_axis = [0] + reduce_axis
@@ -208,9 +210,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                         If `include_background=False`, the weight should not include
                         the background category class 0."""
                     )
-            if self.class_weight.min() < 0:
-                raise ValueError("the value/values of the `weight` should be no less than 0.")
-            # apply class_weight to loss
+            # apply class_weight to loss (weight values validated in __init__)
             f = f * self.class_weight.to(f)
 
         if self.reduction == LossReduction.MEAN.value:
@@ -431,7 +431,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})")
 
         # reducing only spatial dimensions (not batch nor channels)
-        reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
+        reduce_axis: list[int] = list(range(2, len(input.shape)))
         if self.batch:
             reduce_axis = [0] + reduce_axis
 

monai/losses/focal_loss.py

@@ -122,6 +122,8 @@ def __init__(
         else:
             self.alpha = torch.as_tensor(alpha)
         weight = torch.as_tensor(weight) if weight is not None else None
+        if weight is not None and weight.min() < 0:
+            raise ValueError("the value/values of the `weight` should be no less than 0.")
         self.register_buffer("class_weight", weight)
         self.class_weight: None | torch.Tensor
 
@@ -188,9 +190,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
                         If `include_background=False`, the weight should not include
                         the background category class 0."""
                     )
-            if self.class_weight.min() < 0:
-                raise ValueError("the value/values of the `weight` should be no less than 0.")
-            # apply class_weight to loss
+            # apply class_weight to loss (weight values validated in __init__)
             self.class_weight = self.class_weight.to(loss)
             broadcast_dims = [-1] + [1] * len(target.shape[2:])
             self.class_weight = self.class_weight.view(broadcast_dims)

monai/losses/hausdorff_loss.py

@@ -190,7 +190,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             distance = pred_dt**self.alpha + target_dt**self.alpha
 
             running_f = pred_error * distance.to(device)
-            reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
+            reduce_axis: list[int] = list(range(2, len(input.shape)))
             if self.batch:
                 # reducing spatial dimensions and batch
                 reduce_axis = [0] + reduce_axis

monai/losses/tversky.py

@@ -143,7 +143,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
             raise AssertionError(f"ground truth has differing shape ({target.shape}) from input ({input.shape})")
 
         # reducing only spatial dimensions (not batch nor channels)
-        reduce_axis: list[int] = torch.arange(2, len(input.shape)).tolist()
+        reduce_axis: list[int] = list(range(2, len(input.shape)))
         if self.batch:
             # reducing spatial dimensions and batch
             reduce_axis = [0] + reduce_axis