fix: Adding back the LGBM EarlyStopping

microsoft · Oct 29, 2024 · 654194d · 654194d
1 parent d43b8c1
commit 654194d
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/flaml/automl/model.py b/flaml/automl/model.py
@@ -1581,6 +1581,8 @@ def _callback(self, start_time, deadline, free_mem_ratio, env) -> None:
         now = time.time()
         if env.iteration == 0:
             self._time_per_iter = now - start_time
+        if now + self._time_per_iter > deadline:
+            raise EarlyStopException(env.iteration, env.evaluation_result_list)
         if psutil is not None:
             mem = psutil.virtual_memory()
             if mem.available / mem.total < free_mem_ratio: