From a809259e719537ac65bfb839db762ffa7afea4c1 Mon Sep 17 00:00:00 2001
From: Pedro Rodriguez <par@meta.com>
Date: Fri, 24 Jan 2025 10:57:20 -0800
Subject: [PATCH] Use load_async flag to not start MP iterator (#33)

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):