fix #3658

hiyouga · May 11, 2024 · 4777efe · 4777efe
1 parent 58c522c
commit 4777efe
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/src/llmtuner/extras/callbacks.py b/src/llmtuner/extras/callbacks.py
@@ -139,13 +139,15 @@ def on_evaluate(self, args: "TrainingArguments", state: "TrainerState", control:
         r"""
         Event called after an evaluation phase.
         """
-        self._close_thread_pool()
+        if not self.do_train:
+            self._close_thread_pool()
 
     def on_predict(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""
         Event called after a successful prediction.
         """
-        self._close_thread_pool()
+        if not self.do_train:
+            self._close_thread_pool()
 
     def on_log(self, args: "TrainingArguments", state: "TrainerState", control: "TrainerControl", **kwargs):
         r"""