Merge pull request #26 from htem/rhoadesScholar-DataPipe-RandomLocation

Update BaseDataPipe.py
htem · Sep 18, 2023 · 33861b1 · 33861b1
2 parents 94cf386 + 3a65557
commit 33861b1
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/raygun/io/BaseDataPipe.py b/src/raygun/io/BaseDataPipe.py
@@ -28,9 +28,10 @@ def get_source(self, path, src_names, src_specs=None):
 
     def prenet_pipe(self, mode: str = "train"):
         # Make pre-net datapipe
-        prenet_pipe = self.source + gp.RandomLocation()
+        prenet_pipe = self.source
         if mode == "train":
             sections = [
+                gp.RandomLocation(),
                 "reject",
                 "resample",
                 "preprocess",
@@ -41,7 +42,7 @@ def prenet_pipe(self, mode: str = "train"):
         elif mode == "predict":
             sections = ["reject", "resample", "preprocess", "unsqueeze", "stack"]
         elif mode == "test":
-            sections = ["reject", "resample", "preprocess", "unsqueeze", gp.Stack(1)]
+            sections = [gp.RandomLocation(), "reject", "resample", "preprocess", "unsqueeze", gp.Stack(1)]
         else:
             raise ValueError(f"mode={mode} not implemented.")