fix bug: scalar alpha still passed despite ignored warning

Signed-off-by: ytl0623 <david89062388@gmail.com>

Author: ytl0623

monai/losses/focal_loss.py

@@ -164,16 +164,17 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         loss: Optional[torch.Tensor] = None
         input = input.float()
         target = target.float()
-
+        alpha_arg = self.alpha
         if self.use_softmax:
             if not self.include_background and self.alpha is not None:
                 if isinstance(self.alpha, (float, int)):
+                    alpha_arg = None
                     warnings.warn(
                         "`include_background=False`, scalar `alpha` ignored when using softmax.", stacklevel=2
                     )
-            loss = softmax_focal_loss(input, target, self.gamma, self.alpha)
+            loss = softmax_focal_loss(input, target, self.gamma, alpha_arg)
         else:
-            loss = sigmoid_focal_loss(input, target, self.gamma, self.alpha)
+            loss = sigmoid_focal_loss(input, target, self.gamma, alpha_arg)
 
         num_of_classes = target.shape[1]
         if self.class_weight is not None and num_of_classes != 1:

tests/losses/test_focal_loss.py

@@ -79,10 +79,10 @@
 
 TEST_ALPHA_BROADCASTING = []
 for case in TEST_DEVICES:
-    device = case[0]
+    dev = case[0]
     for include_background in [True, False]:
         for use_softmax in [True, False]:
-            TEST_ALPHA_BROADCASTING.append([device, include_background, use_softmax])
+            TEST_ALPHA_BROADCASTING.append([dev, include_background, use_softmax])
 
 
 class TestFocalLoss(unittest.TestCase):