From 57695c6b2c20cf5b64fe11d328b31d88c211f8e1 Mon Sep 17 00:00:00 2001
From: Nathan Chandrasekar <322284620+vera-rubin@users.noreply.github.com>
Date: Sat, 17 Oct 2026 20:41:16 +0000
Subject: [PATCH] Fix DiceLoss ignoring weights for a single foreground class

Check the original prediction channel count so class weights still apply
when excluding background leaves one foreground class. Add regression
coverage for zero and doubled foreground weights.

Assisted-by: Codex / gpt-6.1-sol
Signed-off-by: Nathan Chandrasekar <322284620+vera-rubin@users.noreply.github.com>
---
 monai/losses/dice.py           | 2 +-
 tests/losses/test_dice_loss.py | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/monai/losses/dice.py b/monai/losses/dice.py
index 5561012f0ee..5f8421f3841 100644
--- a/monai/losses/dice.py
+++ b/monai/losses/dice.py
@@ -213,7 +213,7 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         f: torch.Tensor = 1 - numerator / denominator
 
         num_of_classes = target.shape[1]
-        if self.class_weight is not None and num_of_classes != 1:
+        if self.class_weight is not None and n_pred_ch != 1:
             # make sure the lengths of weights are equal to the number of classes
             if self.class_weight.ndim == 0:
                 self.class_weight = torch.as_tensor([self.class_weight] * num_of_classes)
diff --git a/tests/losses/test_dice_loss.py b/tests/losses/test_dice_loss.py
index d8fb5e11957..32b6ba6db0f 100644
--- a/tests/losses/test_dice_loss.py
+++ b/tests/losses/test_dice_loss.py
@@ -184,6 +184,13 @@ def test_shape(self, input_param, input_data, expected_val):
         result = DiceLoss(**input_param).forward(**input_data)
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, rtol=1e-5)
 
+    @parameterized.expand([([0], 0.0), ([2], 0.999995)])
+    def test_single_foreground_weight(self, weight, expected_val):
+        prediction = torch.full((1, 2, 1, 2), 0.5)
+        target = torch.tensor([[[[1.0, 0.0]], [[0.0, 1.0]]]])
+        result = DiceLoss(include_background=False, weight=weight)(prediction, target)
+        np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, rtol=1e-5)
+
     def test_ill_shape(self):
         loss = DiceLoss()
         with self.assertRaisesRegex(AssertionError, ""):