more

mosaicml · Apr 12, 2024 · 6688615 · 6688615
1 parent a088fbf
commit 6688615
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llmfoundry/models/utils/param_init_fns.py b/llmfoundry/models/utils/param_init_fns.py
@@ -328,6 +328,8 @@ def moe_init(
     print(type(module))
     print(isinstance(module, GLU))
     print(isinstance(module, MLP))
+    print(GLU)
+    print(MLP)
     if megablocks is not None and isinstance(module, (
             megablocks.layers.moe.MoE,
             megablocks.layers.dmoe.dMoE,