Spaces:
Running
on
Zero
Running
on
Zero
Use load_async flag to not start MP iterator (#33)
Browse files- bytelatent/args.py +7 -5
bytelatent/args.py
CHANGED
@@ -150,11 +150,13 @@ class DataloaderArgs(BaseModel):
|
|
150 |
enable_byte_ngrams=self.enable_byte_ngrams,
|
151 |
)
|
152 |
packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args)
|
153 |
-
|
154 |
-
|
155 |
-
|
156 |
-
|
157 |
-
|
|
|
|
|
158 |
|
159 |
|
160 |
class TrainArgs(BaseModel):
|
|
|
150 |
enable_byte_ngrams=self.enable_byte_ngrams,
|
151 |
)
|
152 |
packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args)
|
153 |
+
if self.load_async:
|
154 |
+
mp_iterator = MultiprocessIterator(
|
155 |
+
packing_iterator, n_batches_to_prefetch=self.prefetch_size
|
156 |
+
)
|
157 |
+
return mp_iterator
|
158 |
+
else:
|
159 |
+
return packing_iterator
|
160 |
|
161 |
|
162 |
class TrainArgs(BaseModel):
|