diff --git a/tests/models/layers/test_dmoe.py b/tests/models/layers/test_dmoe.py
index 58d110f9c4..e1bfd97939 100644
--- a/tests/models/layers/test_dmoe.py
+++ b/tests/models/layers/test_dmoe.py
@@ -217,12 +217,12 @@ def delete_transformers_cache():
 def test_fwd_equal_dmoe(seqlen: int, precision: str, mlp_type: str):
     delete_transformers_cache()
 
-    import llmfoundry
-    print(llmfoundry.layers_registry.ffns.get_all())
+    from llmfoundry.layers_registry import module_init_fns
+    print(module_init_fns.get_all())
 
     from llmfoundry.models.layers.ffn import resolve_ffn_act_fn  # type: ignore
 
-    print(llmfoundry.layers_registry.ffns.get_all())
+    print(module_init_fns.get_all())
 
     mb_dmoe_config = MPTConfig(d_model=1024,
                                n_heads=32,