Update BaseDataPipe.py

make only train and test pipes include RandomLocation
htem · Sep 5, 2023 · 3a65557 · 3a65557
1 parent 94cf386
commit 3a65557
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.")