fix

mosaicml · Apr 12, 2024 · ce04066 · ce04066
1 parent 3ecd51e
commit ce04066
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -73,6 +73,7 @@ def __init__(
         super().__init__()
 
         self.ffn = build_ffn(
+            name=ffn_config['ffn_type'],
             d_model=d_model,
             expansion_ratio=expansion_ratio,
             device=device,