Function pytorch_adapter

dali/test/python/test_dltensor_operator.py:123–132

(fun, in1, in2)

def pytorch_adapter(fun, in1, in2):
    with torch.cuda.stream(torch_stream):
        tin1 = [torch_dlpack.from_dlpack(dltensor) for dltensor in in1]
        tin2 = [torch_dlpack.from_dlpack(dltensor) for dltensor in in2]
        tout1, tout2 = fun(tin1, tin2)
        out1, out2 = [torch_dlpack.to_dlpack(tout) for tout in tout1], [
            torch_dlpack.to_dlpack(tout) for tout in tout2
        ]
    torch_stream.synchronize()
    return out1, out2


def pytorch_wrapper(fun):