fix?

mosaicml · Apr 12, 2024 · b61d761 · b61d761
1 parent a407d46
commit b61d761
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -73,8 +73,10 @@ def __init__(
         del kwargs  # unused, just to capture any extra args from the config
         super().__init__()
 
+        ffn_type = ffn_config.pop('ffn_type')
+
         self.ffn = build_ffn(
-            name=ffn_config['ffn_type'],
+            name=ffn_type,
             d_model=d_model,
             expansion_ratio=expansion_ratio,
             device=device,