increase dims

allenai · Oct 11, 2024 · 5fe437a · 5fe437a
1 parent ba95e86
commit 5fe437a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/test/nn/moe_test.py b/src/test/nn/moe_test.py
@@ -10,8 +10,8 @@
 @requires_megablocks
 @pytest.mark.parametrize("moe_type", [MoEType.default, MoEType.dropless])
 def test_moe(moe_type):
-    d_model = 64
-    config = MoEConfig(name=moe_type, hidden_size=256, num_experts=4)
+    d_model = 128
+    config = MoEConfig(name=moe_type, hidden_size=512, num_experts=4)
     moe = config.build(d_model=d_model, init_device="cuda")
 
     # Check num params calculation.