diff --git a/src/mflux/weights/weight_handler.py b/src/mflux/weights/weight_handler.py
index fc485ca..14e6ebe 100644
--- a/src/mflux/weights/weight_handler.py
+++ b/src/mflux/weights/weight_handler.py
@@ -75,10 +75,11 @@ class WeightHandler:
         # Reshape and process the huggingface weights
         if "transformer_blocks" in weights:
             for block in weights["transformer_blocks"]:
-                block["ff"] = {
-                    "linear1": block["ff"]["net"][0]["proj"],
-                    "linear2": block["ff"]["net"][2]
-                }
+                if block.get("ff") is not None:
+                    block["ff"] = {
+                        "linear1": block["ff"]["net"][0]["proj"],
+                        "linear2": block["ff"]["net"][2]
+                    }
                 if block.get("ff_context") is not None:
                     block["ff_context"] = {
                         "linear1": block["ff_context"]["net"][0]["proj"],