docarray/tests/unit/array/test_dataloader.py at docarray-v1-fix · docarray/docarray

This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Repository navigation

import time

import numpy as np

from docarray import DocumentArray

def cpu_job(da):

time.sleep(2)

assert da.tensors.shape

return da

def gpu_job(da):

time.sleep(1)

assert da.tensors.shape

def test_dataloader():

N = 100

da = DocumentArray.empty(N)

da.tensors = np.array(255 * np.random.random([N, 32, 32, 3]), dtype=np.uint8)

da.save_binary('da.protobuf.gz')

for da in DocumentArray.dataloader(

'da.protobuf.gz', func=cpu_job, batch_size=64, num_worker=4

gpu_job(da)