From 906b8b314a6de778c5c54a9651c527b1c1c016df Mon Sep 17 00:00:00 2001
From: Grain Team <grain-dev@google.com>
Date: Fri, 16 Jan 2026 09:42:16 -0800
Subject: [PATCH] Internal

PiperOrigin-RevId: 857197203
---
 .../_src/python/dataset/transformations/process_prefetch.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/grain/_src/python/dataset/transformations/process_prefetch.py b/grain/_src/python/dataset/transformations/process_prefetch.py
index d5552dfc..4b271ff1 100644
--- a/grain/_src/python/dataset/transformations/process_prefetch.py
+++ b/grain/_src/python/dataset/transformations/process_prefetch.py
@@ -616,7 +616,11 @@ def multiprocess_prefetch(
     shards.append(worker_ds)
 
   ds = interleave.InterleaveIterDataset(
-      shards, cycle_length=num_workers, iter_buffer_size=buffer_size
+      shards,
+      cycle_length=num_workers,
+      num_make_iter_threads=num_workers,
+      make_iter_buffer_size=num_workers,
+      iter_buffer_size=buffer_size,
   )
   # Apply options from parent dataset because interleave dataset does not
   # propagate options.