Skip to content

Commit 2e5f9c2

Browse files
committed
fix(networks): support QuickNAT without optional SE blocks
Signed-off-by: kyinhub <kevinpyin@gmail.com>
1 parent 3ee058b commit 2e5f9c2

2 files changed

Lines changed: 18 additions & 1 deletion

File tree

monai/networks/nets/quicknat.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,16 @@ class SkipConnectionWithIdx(SkipConnection):
4343
"""
4444

4545
def forward(self, input, indices): # type: ignore[override]
46-
return super().forward(input), indices
46+
submodule_output, _ = self.submodule(input, None)
47+
if self.mode == "cat":
48+
output = torch.cat([input, submodule_output], dim=self.dim)
49+
elif self.mode == "add":
50+
output = torch.add(input, submodule_output)
51+
elif self.mode == "mul":
52+
output = torch.mul(input, submodule_output)
53+
else:
54+
raise NotImplementedError(f"Unsupported mode {self.mode}.")
55+
return output, indices
4756

4857

4958
class SequentialWithIdx(nn.Sequential):

tests/networks/nets/test_quicknat.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,14 @@
3636
]
3737

3838

39+
class TestQuicknatCore(unittest.TestCase):
40+
def test_forward_without_optional_se_dependency(self):
41+
net = Quicknat(num_classes=2, num_channels=1, num_filters=4, se_block=None)
42+
with eval_mode(net):
43+
result = net(torch.randn(1, 1, 32, 32))
44+
self.assertEqual(result.shape, (1, 2, 32, 32))
45+
46+
3947
@unittest.skipUnless(has_se, "squeeze_and_excitation not installed")
4048
class TestQuicknat(unittest.TestCase):
4149
@parameterized.expand(TEST_CASES)

0 commit comments

Comments
 (0)