Function get_device_batch

tests/data/test_data_iterator.py:97–103

(_worker_its)

                global_batch_size=global_batch_size, seed=32, start_index=start)))

        def get_device_batch(_worker_its):
            while True:
                for it in _worker_its:
                    batch = []
                    for _ in range(device_batch_size):
                        batch.append(next(it))
                    yield batch

        device_iterators.append(get_device_batch(worker_iterators))
    torch.utils.data.get_worker_info = bk