fix

ztxz16 · Jul 25, 2024 · 9ad5f65 · 9ad5f65
1 parent 9d79e50
commit 9ad5f65
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tools/fastllm_pytools/llm.py b/tools/fastllm_pytools/llm.py
@@ -687,6 +687,10 @@ def token_healing(self,
                 softmax(out)
                 if (idx < len(real_input)):
                     cur_prob += math.log(out[real_input[idx]])
+                while True:
+                    if (ret <= -1):
+                        break
+                    ret = fastllm_lib.fetch_response_logits_llm_model(self.model, handle, array)
             max_id = -1
             for i in cur_set:
                 if max_id == -1 or out[i] > out[max_id]: