From 4777efe517c05a599f0ccdf9ccf760f3986d126e Mon Sep 17 00:00:00 2001
From: hiyouga <hiyouga@buaa.edu.cn>
Date: Sun, 12 May 2024 01:25:16 +0800
Subject: [PATCH] fix #3658

---
 src/llmtuner/extras/callbacks.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/llmtuner/extras/callbacks.py b/src/llmtuner/extras/callbacks.py
index a142928a..6d24b244 100644
--- a/src/llmtuner/extras/callbacks.py
+++ b/src/llmtuner/extras/callbacks.py
@@ -139,13 +139,15 @@ class LogCallback(TrainerCallback):
         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"""