From 2eac2f69554506e42fffdc28912b6ee13d26dee5 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Tue, 9 Jan 2024 12:41:54 -0800 Subject: [PATCH] Fiddling with iterator wrapping for HF ds streaming --- timm/data/readers/reader_hfids.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/timm/data/readers/reader_hfids.py b/timm/data/readers/reader_hfids.py index bb5fddd7..9f7ce76d 100644 --- a/timm/data/readers/reader_hfids.py +++ b/timm/data/readers/reader_hfids.py @@ -167,9 +167,11 @@ class ReaderHfids(Reader): target_sample_count = self._num_samples_per_worker() sample_count = 0 - ds_iter = iter(self.ds) + if self.is_training: - ds_iter = chain.from_iterable(repeat(ds_iter)) + ds_iter = chain.from_iterable(repeat(self.ds)) + else: + ds_iter = iter(self.ds) for sample in ds_iter: input_data: Image.Image = sample[self.input_key] if self.input_img_mode and input_data.mode != self.input_img_mode: