Update model.py

src/models/model.py

@@ -101,3 +101,11 @@ class PLDataModule(LightningDataModule):
         :param source_max_token_len:
         :param target_max_token_len:
         :param batch_size:
+        :param split:
+        """
+        super().__init__()
+        self.train_df = train_df
+        self.test_df = test_df
+        self.split = split
+        self.batch_size = batch_size
+        self.target_max_token_len = target_max_token_len