fix: CI errors

Signed-off-by: Rusheel Sharma <rusheelhere@gmail.com>

Author: Rusheel86

monai/losses/unified_focal_loss.py

@@ -160,15 +160,15 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         if y_true.shape != y_pred.shape:
             raise ValueError(f"ground truth has different shape ({y_true.shape}) from input ({y_pred.shape})")
 
+        mask = create_ignore_mask(original_y_true if original_y_true is not None else y_true, self.ignore_index)
+        if mask is not None:
+            mask = mask.expand_as(y_true)
+
         y_pred = torch.clamp(y_pred, self.epsilon, 1.0 - self.epsilon)
         cross_entropy = -y_true * torch.log(y_pred)
 
-        # Apply mask from original labels if available
-        mask = create_ignore_mask(original_y_true if original_y_true is not None else y_true, self.ignore_index)
         if mask is not None:
-            cross_entropy = cross_entropy * mask.expand_as(cross_entropy)
-            # Rename for compatibility with reduction block below
-            spatial_mask = mask
+            cross_entropy = cross_entropy * mask
 
         back_ce = torch.pow(1 - y_pred[:, 0], self.gamma) * cross_entropy[:, 0]
         back_ce = (1 - self.delta) * back_ce
@@ -179,11 +179,9 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         loss = torch.stack([back_ce, fore_ce], dim=1)  # [B, 2, H, W]
 
         if self.reduction == LossReduction.MEAN.value:
-            if self.ignore_index is not None and spatial_mask is not None:
-                # Apply mask to loss, then average over valid elements only
-                # loss has shape [B, 2, H, W], spatial_mask has shape [B, 1, H, W]
-                masked_loss = loss * spatial_mask.expand_as(loss)
-                return masked_loss.sum() / (spatial_mask.expand_as(loss).sum().clamp(min=1e-5))
+            if mask is not None:
+                masked_loss = loss * mask
+                return masked_loss.sum() / mask.expand_as(loss).sum().clamp(min=1e-5)
             return loss.mean()
         if self.reduction == LossReduction.SUM.value:
             return loss.sum()