Merge pull request #19 from antimatter15/patch-1

Fix LoRa weight merging in export It can't hurt
tloen · Mar 16, 2023 · 6681523 · 6681523
2 parents b8c32be + dde8995
commit 6681523
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/export_state_dict_checkpoint.py b/export_state_dict_checkpoint.py
@@ -21,7 +21,12 @@
     torch_dtype=torch.float16,
 )
 
-lora_model.eval()  # merge weights
+# merge weights
+for layer in lora_model.base_model.model.model.layers:
+    layer.self_attn.q_proj.merge_weights = True
+    layer.self_attn.v_proj.merge_weights = True
+
+lora_model.train(False)
 
 lora_model_sd = lora_model.state_dict()