Use load_async flag to not start MP iterator (#33)
Some checks failed
Lint with Black / lint (push) Has been cancelled
Lint with isort / lint (push) Has been cancelled

Summary:

Test Plan:
This commit is contained in:
Pedro Rodriguez 2025-01-24 10:57:20 -08:00 committed by GitHub
parent bc42cebd7d
commit a809259e71
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -150,11 +150,13 @@ class DataloaderArgs(BaseModel):
enable_byte_ngrams=self.enable_byte_ngrams, enable_byte_ngrams=self.enable_byte_ngrams,
) )
packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args) packing_iterator = PackingIterator(sampling_iterator, packing_args=packing_args)
mp_iterator = MultiprocessIterator( if self.load_async:
packing_iterator, n_batches_to_prefetch=self.prefetch_size mp_iterator = MultiprocessIterator(
) packing_iterator, n_batches_to_prefetch=self.prefetch_size
)
return mp_iterator return mp_iterator
else:
return packing_iterator
class TrainArgs(BaseModel): class TrainArgs(BaseModel):