From 550b925af2281dfe414b75a42589bc6cbad7c725 Mon Sep 17 00:00:00 2001
From: Galunid <karolek1231456@gmail.com>
Date: Sun, 29 Oct 2023 02:06:41 +0100
Subject: [PATCH] Missing variable

---
 model.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/model.py b/model.py
index 66b350f65..c35fe3bbf 100644
--- a/model.py
+++ b/model.py
@@ -427,6 +427,12 @@ class BaichuanModel(Model):
         head_count = self.hparams["num_attention_heads"]
         tensor_map = gguf.get_tensor_name_map(self.model_arch, block_count)
 
+        if "num_key_value_heads" in self.hparams:
+            head_count_kv = self.hparams["num_key_value_heads"]
+        else:
+            head_count_kv = head_count
+
+
         for i in range(block_count):
             if f"model.layers.{i}.self_attn.W_pack.weight" in model_kv:
                 print(f"Unpacking and permuting layer {i}")