MCPcopy Create free account
hub / github.com/FreeformRobotics/OTS / __next__

Method __next__

lib/utils/data/dataloader.py:261–286  ·  view source on GitHub ↗
(self)

Source from the content-addressed store, hash-verified

259 return self.data_queue.get()
260
261 def __next__(self):
262 if self.num_workers == 0: # same-process loading
263 indices = next(self.sample_iter) # may raise StopIteration
264 batch = self.collate_fn([self.dataset[i] for i in indices])
265 if self.pin_memory:
266 batch = pin_memory_batch(batch)
267 return batch
268
269 # check if the next sample has already been generated
270 if self.rcvd_idx in self.reorder_dict:
271 batch = self.reorder_dict.pop(self.rcvd_idx)
272 return self._process_next_batch(batch)
273
274 if self.batches_outstanding == 0:
275 self._shutdown_workers()
276 raise StopIteration
277
278 while True:
279 assert (not self.shutdown and self.batches_outstanding > 0)
280 idx, batch = self._get_batch()
281 self.batches_outstanding -= 1
282 if idx != self.rcvd_idx:
283 # store out-of-order samples
284 self.reorder_dict[idx] = batch
285 continue
286 return self._process_next_batch(batch)
287
288 next = __next__ # Python 2 compatibility
289

Callers

nothing calls this directly

Calls 4

_process_next_batchMethod · 0.95
_shutdown_workersMethod · 0.95
_get_batchMethod · 0.95
pin_memory_batchFunction · 0.85

Tested by

no test coverage detected