Allow non train mode for wds reader to operate w/o sample count, exhaust iterator

pull/2319/head
Ross Wightman 2024-10-31 15:39:42 -07:00 committed by Ross Wightman
parent f689c850b9
commit f5b58e31a2
1 changed files with 2 additions and 2 deletions

View File

@ -304,8 +304,8 @@ class ReaderWds(Reader):
self.num_samples = num_samples
else:
self.num_samples = self.split_info.num_samples
if not self.num_samples:
raise RuntimeError(f'Invalid split definition, num_samples not specified.')
if is_training and not self.num_samples:
raise RuntimeError(f'Invalid split definition, num_samples not specified in train mode.')
self.remap_class = False
if class_map:
self.class_to_idx = load_class_map(class_map)