Update model.py

src/models/model.py

@@ -121,3 +121,9 @@ class PLDataModule(LightningDataModule):
         )
         self.test_dataset = DataModule(
             self.test_df,
+            self.tokenizer,
+            self.source_max_token_len,
+            self.target_max_token_len,
+        )
+
+    def train_dataloader(self):