fix

mosaicml · Apr 12, 2024 · d110f74 · d110f74
1 parent bc87d1d
commit d110f74
Showing 1 changed file with 3 additions and 3 deletions.
diff --git 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,