diff --git a/src/deepmol/pipeline_optimization/objective_wrapper.py b/src/deepmol/pipeline_optimization/objective_wrapper.py
index 652c3746..1d5c110a 100644
--- a/src/deepmol/pipeline_optimization/objective_wrapper.py
+++ b/src/deepmol/pipeline_optimization/objective_wrapper.py
@@ -42,8 +42,8 @@ def __init__(self, objective_steps, study, direction, save_top_n, trial_timeout=
     def run_with_timeout_windows(self):
 
         @timeout(self.trial_timeout)
-        def run_with_timeout():
-            return self._run()
+        def run_with_timeout(trial):
+            return self._run(trial)
         
         return run_with_timeout