par-meta commited on
Commit
a809259
·
unverified ·
1 Parent(s): bc42ceb

Use load_async flag to not start MP iterator (#33)

Browse files
Files changed (1) hide show
  1. 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
- mp_iterator = MultiprocessIterator(
154
- packing_iterator, n_batches_to_prefetch=self.prefetch_size
155
- )
156
-
157
- return mp_iterator
 
 
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):