reinitialize trainer after fitting

basf · Dec 4, 2024 · 86d946f · 86d946f
1 parent 2f9baee
commit 86d946f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/molpipeline/estimators/chemprop/abstract.py b/molpipeline/estimators/chemprop/abstract.py
@@ -100,9 +100,7 @@ def __init__(
         self.trainer_params = get_params_trainer(self.lightning_trainer)
         self.set_params(**kwargs)
 
-    def _update_trainer(
-        self,
-    ) -> None:
+    def _update_trainer(self) -> None:
         """Update the trainer for the model."""
         trainer_params = dict(self.trainer_params)
         if self.model_ckpoint_params:
@@ -139,6 +137,8 @@ def fit(
             X, batch_size=self.batch_size, num_workers=self.n_jobs
         )
         self.lightning_trainer.fit(self.model, training_data)
+        # The trainer is reinitalized to avoid storing the training data
+        self._update_trainer()
         return self
 
     def set_params(self, **params: Any) -> Self: