From bd461af91a86b572442fb4170c288603dbd5e6e5 Mon Sep 17 00:00:00 2001
From: Pedro Rodriguez <me@pedro.ai>
Date: Fri, 24 Jan 2025 18:56:18 +0000
Subject: [PATCH] Use load_async flag to not start MP iterator

Summary:

Test Plan:
---
 bytelatent/args.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/bytelatent/args.py b/bytelatent/args.py
index b9144c6..a332c89 100644
--- a/bytelatent/args.py
+++ b/bytelatent/args.py
@@ -150,11 +150,13 @@ class DataloaderArgs(BaseModel):
             enable_byte_ngrams=self.enable_byte_ngrams,
         )
         packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args)
-        mp_iterator = MultiprocessIterator(
-            packing_iterator, n_batches_to_prefetch=self.prefetch_size
-        )
-
-        return mp_iterator
+        if self.load_async:
+            mp_iterator = MultiprocessIterator(
+                packing_iterator, n_batches_to_prefetch=self.prefetch_size
+            )
+            return mp_iterator
+        else:
+            return packing_iterator
 
 
 class TrainArgs(BaseModel):