From 7d27f1c4c1e1d25ab5326618d509f6a162c2e3b1 Mon Sep 17 00:00:00 2001 From: filipstrand Date: Thu, 12 Sep 2024 20:43:06 +0200 Subject: [PATCH] Adjust FF block only if it exists --- src/mflux/weights/weight_handler.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) 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"],