def __next__(self):
index = self._next_index() # may raise StopIteration
data = self.dataset_fetcher.fetch(index) # may raise StopIteration
if self.dataset._should_yield_raw_example:
index = [idx[0] for idx in index]
examples, data = data
self.dataset._add_cached_examples(index, examples)
if self.pin_memory:
data = _pin_memory(data)
return data
def __next__(self):
if self.num_workers == 0: # same-process loading
indices = next(self.sample_iter) # may raise StopIteration
batch = self.collate_fn([self.dataset[i] for i in indices])
if self.dataset._should_yield_raw_example:
indices = [index[0] for index in indices]
examples, batch = batch
self.dataset._add_cached_examples(indices, examples)
if self.pin_memory:
batch = _pin_memory(batch)
else:
batch = super().__next__()
if (self._batch_size is not None
and batch.batch_size < self.dataset.batch_size
and not self.dataset.hparams.allow_smaller_final_batch):
raise StopIteration
return batch
