adding num workers to mock img ds

ben_publication/mock_dm.py

@@ -17,7 +17,7 @@ class MockDataset(Dataset):
 
 
 class MockDataModule(LightningDataModule):
-    def __init__(self, dims, clss, train_length, val_length, test_length, bs):
+    def __init__(self, dims, clss, train_length, val_length, test_length, bs, num_workers=8):
         super().__init__()
         self.dims = dims
         self.clss = clss
@@ -25,16 +25,17 @@ class MockDataModule(LightningDataModule):
         self.val_length = val_length
         self.test_length = test_length
         self.batch_size = bs
+        self.num_workers = num_workers
         self.train_ds = MockDataset(self.dims, self.clss, self.length)
         self.val_ds = MockDataset(self.dims, self.clss, self.val_length)
         self.test_ds = MockDataset(self.dims, self.clss, self.test_length)
         print("MockDataModule initialized")
 
     def train_dataloader(self):
-        return torch.utils.data.DataLoader(self.train_ds, batch_size=self.batch_size)
+        return torch.utils.data.DataLoader(self.train_ds, batch_size=self.batch_size, num_workers=self.num_workers)
 
     def val_dataloader(self):
-        return torch.utils.data.DataLoader(self.val_ds, batch_size=self.batch_size)
+        return torch.utils.data.DataLoader(self.val_ds, batch_size=self.batch_size, num_workers=self.num_workers)
 
     def test_dataloader(self):
-        return torch.utils.data.DataLoader(self.test_ds, batch_size=self.batch_size)
+        return torch.utils.data.DataLoader(self.test_ds, batch_size=self.batch_size, num_workers=self.num_workers)

scripts/train_BENv2.py

@@ -124,6 +124,7 @@ def main(
             val_length=120_000,
             test_length=120_000,
             bs=bs,
+            num_workers=workers,
         )
     else:
         dm = BENv2DataModule(