Fix test input shape and class name for UMambaUNet

Author: billyluu5704

networks/nets/__init__.py

@@ -144,4 +144,3 @@
 from .vnet import VNet
 from .voxelmorph import VoxelMorph, VoxelMorphUNet
 from .vqvae import VQVAE
-from .u_mamba import UMamba

networks/nets/u_mamba.py

@@ -1,22 +1,18 @@
 import unittest
 import torch
-from monai.networks.nets import UMamba
+from monai.networks.nets import UMambaUNet
 
 class TestUMamba(unittest.TestCase):
     def test_forward_shape(self):
         # Set up input dimensions and model
-        input_tensor = torch.randn(2, 1, 64, 64)  # (batch_size, channels, H, W)
-        model = UMamba(in_channels=1, out_channels=2)  # example args
-
-        # Forward pass
+        input_tensor = torch.randn(2, 1, 16, 64, 64)
+        model = UMambaUNet(in_channels=1, out_channels=2)
         output = model(input_tensor)
-
-        # Assert output shape matches expectation
-        self.assertEqual(output.shape, (2, 2, 64, 64))  # adjust if necessary
+        self.assertEqual(output.shape, (2, 2, 16, 64, 64))
 
     def test_script(self):
         # Test JIT scripting if supported
-        model = UMamba(in_channels=1, out_channels=2)
+        model = UMambaUNet(in_channels=1, out_channels=2)
         scripted = torch.jit.script(model)
         x = torch.randn(1, 1, 64, 64)
         out = scripted(x)