fix bug where val_ds still has augmentations turned on

talmolab · Aug 30, 2024 · 5fb8256 · 5fb8256
1 parent 1680ab7
commit 5fb8256
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/dreem/datasets/tracking_dataset.py b/dreem/datasets/tracking_dataset.py
@@ -99,6 +99,7 @@ def make_train_splits(self, splits: tuple[float]) -> None:
 
  if train_frac != 0:
  dataset_to_split = copy(train_ds)
+ dataset_to_split.augmentations = None
  dataset_size = len(dataset_to_split)
  train_inds, val_inds = train_test_split(
  range(dataset_size),
@@ -117,12 +118,10 @@ def make_train_splits(self, splits: tuple[float]) -> None:
 
  if len(test_inds) != 0:
  test_ds = Subset(dataset_to_split, test_inds)
- test_ds.dataset.augmentations = None
  else:
  test_ds = None
 
  val_ds = Subset(dataset_to_split, val_inds)
- val_ds.dataset.augmentations = None
 
  train_ds = Subset(train_ds, train_inds)