fix

intel-analytics · Jul 1, 2024 · be7cc05 · be7cc05
1 parent d1af961
commit be7cc05
Showing 1 changed file with 3 additions 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
@@ -611,7 +611,9 @@ async def process_step(self, tokenizer, result_dict):
 
             if (cur_batch is not None) and (not cur_batch.stopped) and (cur_input is None):
                 cur_id = cur_batch.batch_id
-                cur_batch = self.prepare_batch(cur_batch)
+                # cur_batch = self.prepare_batch(cur_batch)
+                if cur_batch.prefilled_index >= cur_batch.batch_size:
+                    cur_batch.partial_prefilling = 0
                 if cur_batch.partial_prefilling > 0:
                     next_ids = torch.empty((cur_batch.partial_prefilling, 1,), device=f'xpu:{self.rank}', dtype=torch.int64)
                 else: