Update model.py

src/models/model.py

@@ -13,3 +13,5 @@ import pytorch_lightning as pl
 from pytorch_lightning.loggers import MLFlowLogger
 from pytorch_lightning import Trainer
 from pytorch_lightning.callbacks.early_stopping import EarlyStopping
+from pytorch_lightning import LightningDataModule
+from pytorch_lightning import LightningModule