Gagan Bhatia commited on
Commit
68248ca
1 Parent(s): f9e196a

Update model.py

Browse files
Files changed (1) hide show
  1. src/models/model.py +5 -0
src/models/model.py CHANGED
@@ -131,3 +131,8 @@ class PLDataModule(LightningDataModule):
131
  return DataLoader(
132
  self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=2
133
  )
 
 
 
 
 
 
131
  return DataLoader(
132
  self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=2
133
  )
134
+
135
+ def test_dataloader(self):
136
+ """ test dataloader """
137
+ return DataLoader(
138
+ self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=2