reinitialize trainer after fitting (#112)

Author: c-w-feldmann

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: