Replace prefetch with mp_prefetch in PyGrain.

This is a no-op change. PiperOrigin-RevId: 689791881
google-research · Oct 30, 2024 · 35efcd5 · 35efcd5
1 parent bc8767c
commit 35efcd5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/kauldron/data/py/base.py b/kauldron/data/py/base.py
@@ -121,7 +121,7 @@ def _root_ds(self) -> grain.IterDataset:
           num_workers,
           enable_profiling=self.enable_profiling,
       )
-      ds = ds.prefetch(multiprocessing_options)
+      ds = ds.mp_prefetch(multiprocessing_options)
     return ds
 
   def __iter__(self) -> iterators.Iterator: