more swiglu

zachcp · Dec 5, 2024 · 74cf6b3 · 74cf6b3
1 parent b706689
commit 74cf6b3
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/ferritin-esm/src/esm/layers/blocks.rs b/ferritin-esm/src/esm/layers/blocks.rs
@@ -53,7 +53,7 @@ pub struct UnifiedTransformerBlock {
     attn: Option<MultiHeadAttention>,
     use_geom_attn: bool,
     geom_attn: Option<GeometricReasoningOriginalImpl>,
-    ffn: nn::Sequential,
+    ffn: SwiGLU,
     scaling_factor: f64,
 }
 
@@ -170,16 +170,16 @@ impl UnifiedTransformerBlock {
         };
 
         let ffn = match ffn_type {
-            Ffn_Type::GLU => unimplemented!(),
-            Ffn_Type::SWIGLU => Swiglu::load(vb, config),
+            Ffn_Type::SWIGLU => SwiGLU::load(vb, config),
+            _ => unimplemented!(), // Ffn_Type::GLU => unimplemented!(),
         };
 
         Self {
             use_plain_attn,
             attn,
             use_geom_attn,
             geom_attn,
-            ffn,
+            ffn: ffn.unwrap(),
             scaling_factor: residue_scaling_factor,
         }
     }