fix

intel · Jul 4, 2024 · 7e74090 · 7e74090
1 parent fdc87df
commit 7e74090
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python/llm/src/ipex_llm/transformers/pipeline_parallel.py b/python/llm/src/ipex_llm/transformers/pipeline_parallel.py
@@ -564,7 +564,7 @@ def model_step(self, input, cur_batch):
                     _pre_output = torch.cat((_pre_output, tmp_output), dim=0)
                 self.partial_output_dict[cur_id] = _pre_output
         else:
-            _past_key_values = self.update_kv_cache(_past_key_values, cur_id)
+            _past_key_values = self.update_kv_cache(output.past_key_values, cur_id)
             self.past_key_values_dict[cur_id] = _past_key_values
         torch.xpu.synchronize()
         if not self.pp_config.is_tail: