diff --git a/src/openpi/training/data_loader.py b/src/openpi/training/data_loader.py
index e2ee7dd06b..d7f544df97 100644
--- a/src/openpi/training/data_loader.py
+++ b/src/openpi/training/data_loader.py
@@ -424,6 +424,7 @@ def __init__(
                 jax.sharding.PartitionSpec("B"),
             )
         self._num_batches = num_batches
+        self._sampler = sampler
 
         mp_context = None
         if num_workers > 0:
@@ -449,6 +450,15 @@ def __init__(
     def torch_loader(self) -> torch.utils.data.DataLoader:
         return self._data_loader
 
+    def set_epoch(self, epoch: int) -> None:
+        if hasattr(self._sampler, "set_epoch"):
+            self._sampler.set_epoch(epoch)
+
+    def __len__(self) -> int:
+        if self._num_batches is not None:
+            return self._num_batches
+        return len(self._data_loader)
+
     def __iter__(self):
         num_items = 0
         while True:
@@ -535,6 +545,13 @@ def __init__(self, data_config: _config.DataConfig, data_loader: TorchDataLoader
     def data_config(self) -> _config.DataConfig:
         return self._data_config
 
+    def set_epoch(self, epoch: int) -> None:
+        if hasattr(self._data_loader, "set_epoch"):
+            self._data_loader.set_epoch(epoch)
+
+    def __len__(self) -> int:
+        return len(self._data_loader)
+
     def __iter__(self):
         for batch in self._data_loader:
             yield _model.Observation.from_dict(batch), batch["actions"]