Function _test_external_source_callback

dali/test/python/test_external_source_impl.py:239–257

(use_fn_api, batch, as_tensor, device)

def _test_external_source_callback(use_fn_api, batch, as_tensor, device):
    iter_num = 5
    batch_size = 9
    pipe = Pipeline(batch_size, 3, 0)

    # this should produce a single Tensor / TensorList per batch,
    # not wrapped in additional list
    source = TestIterator(iter_num, batch_size, 4, device == "gpu")

    iter_in = iter(source) if batch else iter(SampleIterator(iter(source)))

    if use_fn_api:
        input = fn.external_source(lambda: next(iter_in), device=device, batch=batch)
    else:
        ext_source = ops.ExternalSource(lambda: next(iter_in), device=device, batch=batch)
        input = ext_source()
    pipe.set_outputs(input)

    run_and_check(pipe, source)


def test_external_source_callback():