diff --git a/shimmer/data/dataset.py b/shimmer/data/dataset.py
index 8776ccc..6e3b19f 100644
--- a/shimmer/data/dataset.py
+++ b/shimmer/data/dataset.py
@@ -87,6 +87,7 @@ def __init__(
             self.domains[domain.kind] = domain_cls(
                 dataset_path,
                 split,
+                max_size,
                 transform,
                 self.domain_args.get(domain.kind, None),
             )