Fix QuickNAT indexed skip connections

monai/networks/nets/quicknat.py

@@ -43,7 +43,28 @@ class SkipConnectionWithIdx(SkipConnection):
     """
 
     def forward(self, input, indices):  # type: ignore[override]
-        return super().forward(input), indices
+        """Combine the input with the indexed submodule output.
+
+        Args:
+            input: input tensor for the skip connection and submodule.
+            indices: pooling indices to preserve for the following decoder.
+
+        Returns:
+            A tuple containing the combined tensor and the original indices.
+
+        Raises:
+            NotImplementedError: if the configured skip mode is unsupported.
+        """
+        submodule_output, _ = self.submodule(input, None)
+        if self.mode == "cat":
+            output = torch.cat([input, submodule_output], dim=self.dim)
+        elif self.mode == "add":
+            output = torch.add(input, submodule_output)
+        elif self.mode == "mul":
+            output = torch.mul(input, submodule_output)
+        else:
+            raise NotImplementedError(f"Unsupported mode {self.mode}.")
+        return output, indices
 
 
 class SequentialWithIdx(nn.Sequential):
@@ -325,7 +346,7 @@ class Quicknat(nn.Module):
         stride_convolution: convolution stride. Defaults to 1.
         pool: kernel size of the pooling layer,
         stride_pool: stride for the pooling layer.
-        se_block: Squeeze and Excite block type to be included, defaults to None. Valid options : NONE, CSE, SSE, CSSE,
+        se_block: Squeeze and Excite block type to include. Use ``None`` or ``"None"`` to disable it; valid enabled values are ``"CSE"``, ``"SSE"``, and ``"CSSE"``.
         droup_out: dropout ratio. Defaults to no dropout.
         act: activation type and arguments. Defaults to PReLU.
         norm: feature normalization type and arguments. Defaults to instance norm.
@@ -358,7 +379,7 @@ def __init__(
         pool: int = 2,
         stride_pool: int = 2,
         # Valid options : NONE, CSE, SSE, CSSE
-        se_block: str = "None",
+        se_block: str | None = "None",
         drop_out: float = 0,
         act: tuple | str = Act.PRELU,
         norm: tuple | str = Norm.INSTANCE,

tests/networks/nets/test_quicknat.py

@@ -18,11 +18,21 @@
 
 from monai.networks import eval_mode
 from monai.networks.nets import Quicknat
+from monai.networks.nets.quicknat import SkipConnectionWithIdx
 from monai.utils import optional_import
 from tests.test_utils import test_script_save
 
 _, has_se = optional_import("squeeze_and_excitation")
 
+
+class _DoubleWithIdx(torch.nn.Module):
+    """Double tensor values using QuickNAT's indexed-module signature."""
+
+    def forward(self, input, indices):
+        """Return the doubled tensor and unchanged indices."""
+        return input * 2, indices
+
+
 TEST_CASES = [
     # params, input_shape, expected_shape
     [{"num_classes": 1, "num_channels": 1, "num_filters": 1, "se_block": None}, (1, 1, 32, 32), (1, 1, 32, 32)],
@@ -36,6 +46,38 @@
 ]
 
 
+class TestQuicknatCore(unittest.TestCase):
+    """Test QuickNAT paths that do not require optional dependencies."""
+
+    @parameterized.expand(["cat", "add", "mul"])
+    def test_skip_connection_modes_preserve_indices(self, mode):
+        """Verify each fusion mode and preservation of pooling indices."""
+        skip = SkipConnectionWithIdx(_DoubleWithIdx(), mode=mode)
+        input = torch.tensor([[[[2.0]]]])
+        indices = torch.tensor([[[[3]]]])
+
+        output, returned_indices = skip(input, indices)
+
+        expected = {"cat": torch.cat([input, input * 2], dim=1), "add": input * 3, "mul": input * input * 2}[mode]
+        self.assertTrue(torch.equal(output, expected))
+        self.assertIs(returned_indices, indices)
+
+    def test_skip_connection_rejects_unsupported_mode(self):
+        """Verify unsupported fusion modes fail explicitly."""
+        skip = SkipConnectionWithIdx(_DoubleWithIdx(), mode="cat")
+        skip.mode = "unsupported"
+
+        with self.assertRaisesRegex(NotImplementedError, "Unsupported mode"):
+            skip(torch.ones((1, 1, 1, 1)), torch.ones((1, 1, 1, 1)))
+
+    def test_forward_without_optional_se_dependency(self):
+        """Verify QuickNAT runs when squeeze-and-excitation is disabled."""
+        net = Quicknat(num_classes=2, num_channels=1, num_filters=4, se_block=None)
+        with eval_mode(net):
+            result = net(torch.randn(1, 1, 32, 32))
+        self.assertEqual(result.shape, (1, 2, 32, 32))
+
+
 @unittest.skipUnless(has_se, "squeeze_and_excitation not installed")
 class TestQuicknat(unittest.TestCase):
     @parameterized.expand(TEST_CASES)