Validate imratio and input shape, added test cases for it and fix non-binary target test

Signed-off-by: Shubham Chandravanshi <shubham.chandravanshi378@gmail.com>

Author: shubham-61969

monai/losses/aucm_loss.py

@@ -65,6 +65,7 @@ def __init__(
 
         Raises:
             ValueError: When ``version`` is not one of ["v1", "v2"].
+            ValueError: When ``imratio`` is not in [0, 1].
 
         Example:
             >>> import torch
@@ -77,6 +78,8 @@ def __init__(
         super().__init__(reduction=LossReduction(reduction).value)
         if version not in ["v1", "v2"]:
             raise ValueError(f"version should be 'v1' or 'v2', got {version}")
+        if imratio is not None and not (0.0 <= imratio <= 1.0):
+            raise ValueError(f"imratio must be in [0, 1], got {imratio}")
         self.margin = margin
         self.imratio = imratio
         self.version = version
@@ -95,8 +98,11 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
 
         Raises:
             ValueError: When input or target have incorrect shapes.
+            ValueError: When input or target have fewer than 2 dimensions.
             ValueError: When target contains non-binary values.
         """
+        if input.ndim < 2 or target.ndim < 2:
+            raise ValueError("Input and target must have at least 2 dimensions (B, C, ...)")
         if input.shape[1] != 1:
             raise ValueError(f"Input should have 1 channel for binary classification, got {input.shape[1]}")
         if target.shape[1] != 1:

tests/losses/test_aucm_loss.py

@@ -45,6 +45,13 @@ def test_invalid_version(self):
         with self.assertRaises(ValueError):
             AUCMLoss(version="invalid")
 
+    def test_invalid_imratio(self):
+        """Test that invalid imratio raises ValueError."""
+        with self.assertRaises(ValueError):
+            AUCMLoss(imratio=1.5)
+        with self.assertRaises(ValueError):
+            AUCMLoss(imratio=-0.1)
+
     def test_invalid_input_shape(self):
         """Test that invalid input shape raises ValueError."""
         loss_fn = AUCMLoss()
@@ -61,6 +68,14 @@ def test_invalid_target_shape(self):
         with self.assertRaises(ValueError):
             loss_fn(input, target)
 
+    def test_insufficient_dimensions(self):
+        """Test that tensors with insufficient dimensions raise ValueError."""
+        loss_fn = AUCMLoss()
+        input = torch.randn(32)  # 1D tensor
+        target = torch.randint(0, 2, (32, 1)).float()
+        with self.assertRaises(ValueError):
+            loss_fn(input, target)
+
     def test_shape_mismatch(self):
         """Test that mismatched shapes raise ValueError."""
         loss_fn = AUCMLoss()
@@ -73,7 +88,7 @@ def test_non_binary_target(self):
         """Test that non-binary target values raise ValueError."""
         loss_fn = AUCMLoss()
         input = torch.randn(32, 1)
-        target = torch.tensor([[0.5], [1.0], [2.0]] * 10 + [[0.0]])  # Contains non-binary values
+        target = torch.tensor([[0.5], [1.0], [2.0], [0.0]] * 8)  # 32x1, still non-binary
         with self.assertRaises(ValueError):
             loss_fn(input, target)