LLM: fix chatglm kv cache (#9215)

2023-10-18 19:09:53 +08:00 · 2023-10-18 19:09:53 +08:00 · 942d6418e7
commit 942d6418e7
parent 0765f94770
1 changed files with 2 additions and 0 deletions
--- a/python/llm/src/bigdl/llm/transformers/models/chatglm.py
+++ b/python/llm/src/bigdl/llm/transformers/models/chatglm.py
@ -77,6 +77,8 @@ def attention_fn(
                                                       device=device)
            new_cache_k[:] = cache_k
            new_cache_v[:] = cache_v
+            cache_k = new_cache_k
+            cache_v = new_cache_v
        key_layer, value_layer = append_kv_cache(cache_k, cache_v, key_layer, value_layer)

    elif use_cache: