Update model.py

src/models/model.py

@@ -141,3 +141,12 @@ class PLDataModule(LightningDataModule):
     def val_dataloader(self):
         """ validation dataloader """
         return DataLoader(
+            self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=2
+        )
+
+
+class LightningModel(LightningModule):
+    """ PyTorch Lightning Model class"""
+
+    def __init__(self, tokenizer, model, output: str = "outputs"):
+        """