Remove trainer from DeepAR and TFT

tinkoff-ai · Nov 25, 2022 · 276ec77 · 276ec77
1 parent 45e62b3
commit 276ec77
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 6 deletions.
diff --git a/etna/models/nn/deepar.py b/etna/models/nn/deepar.py
@@ -103,7 +103,6 @@ def __init__(
         self.trainer_kwargs = trainer_kwargs if trainer_kwargs is not None else dict()
         self.quantiles_kwargs = quantiles_kwargs if quantiles_kwargs is not None else dict()
         self.model: Optional[Union[LightningModule, DeepAR]] = None
-        self.trainer: Optional[pl.Trainer] = None
         self._last_train_timestamp = None
         self._freq: Optional[str] = None
 
@@ -163,11 +162,11 @@ def fit(self, ts: TSDataset) -> "DeepARModel":
         )
         trainer_kwargs.update(self.trainer_kwargs)
 
-        self.trainer = pl.Trainer(**trainer_kwargs)
+        trainer = pl.Trainer(**trainer_kwargs)
 
         train_dataloader = pf_transform.pf_dataset_train.to_dataloader(train=True, batch_size=self.batch_size)
 
-        self.trainer.fit(self.model, train_dataloader)
+        trainer.fit(self.model, train_dataloader)
 
         return self
 

diff --git a/etna/models/nn/tft.py b/etna/models/nn/tft.py
@@ -110,7 +110,6 @@ def __init__(
         self.trainer_kwargs = trainer_kwargs if trainer_kwargs is not None else dict()
         self.quantiles_kwargs = quantiles_kwargs if quantiles_kwargs is not None else dict()
         self.model: Optional[Union[LightningModule, TemporalFusionTransformer]] = None
-        self.trainer: Optional[pl.Trainer] = None
         self._last_train_timestamp = None
         self._freq: Optional[str] = None
 
@@ -170,11 +169,11 @@ def fit(self, ts: TSDataset) -> "TFTModel":
         )
         trainer_kwargs.update(self.trainer_kwargs)
 
-        self.trainer = pl.Trainer(**trainer_kwargs)
+        trainer = pl.Trainer(**trainer_kwargs)
 
         train_dataloader = pf_transform.pf_dataset_train.to_dataloader(train=True, batch_size=self.batch_size)
 
-        self.trainer.fit(self.model, train_dataloader)
+        trainer.fit(self.model, train_dataloader)
 
         return self