Merge pull request #29 from kalininalab/preprocess

Preprocess
kalininalab · Aug 16, 2024 · 9376dcc · 9376dcc
2 parents fcf1d1b + a348403
commit 9376dcc
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/src/data.py b/src/data.py
@@ -41,7 +41,7 @@ def __init__(self, data_dir: str | Path, layer_num: int, batch_size: int, num_wo
         self.batch_size = batch_size
         self.num_workers = num_workers
 
-    def setup(self):
+    def setup(self, stage=None):
         self.train = DownstreamDataset(self.data_dir / "train", self.layer_num)
         self.valid = DownstreamDataset(self.data_dir / "valid", self.layer_num)
         self.test = DownstreamDataset(self.data_dir / "test", self.layer_num)

diff --git a/train.py b/train.py
@@ -76,6 +76,7 @@ def train(
 
 
 if __name__ == "__main__":
+    torch.multiprocessing.set_start_method("spawn")  # good solution !!!!
     import jsonargparse
 
     jsonargparse.CLI(train)