diff --git a/tests/unit_tests/transformer/moe/test_paged_stashing.py b/tests/unit_tests/transformer/moe/test_paged_stashing.py
index 04dbbaa5ec6..df78214cb2b 100644
--- a/tests/unit_tests/transformer/moe/test_paged_stashing.py
+++ b/tests/unit_tests/transformer/moe/test_paged_stashing.py
@@ -15,6 +15,7 @@
     paged_stash_init_chunk_handler,
     paged_stash_reset,
 )
+from megatron.core.transformer.spec_utils import get_submodules
 from megatron.core.transformer.transformer_config import TransformerConfig
 from megatron.core.utils import is_te_min_version
 from megatron.training.initialize import _set_random_seed
@@ -136,7 +137,7 @@ def _create_moe_layer(self, layer_number=0):
         quantization_context = get_fp8_context(self.config, layer_number, is_init=True)
         with quantization_context:
             moe_layer = (
-                MoELayer(self.config, transformer_layer_spec.submodules.mlp.submodules)
+                MoELayer(self.config, get_submodules(transformer_layer_spec.submodules.mlp))
                 .cuda()
                 .to(dtype=self.test_dtype)
             )