diff --git a/dreem/datasets/tracking_dataset.py b/dreem/datasets/tracking_dataset.py
index 9bf7f7e..1cf4d28 100644
--- a/dreem/datasets/tracking_dataset.py
+++ b/dreem/datasets/tracking_dataset.py
@@ -131,7 +131,7 @@ def make_train_splits(self, splits: tuple[float]) -> None:
             dataset_size = len(dataset_to_split)
 
             val_inds, test_inds = train_test_split(
-                range(test_size),
+                range(dataset_size),
                 train_size=val_frac,
                 test_size=test_frac,
             )