Function build_and_run

tests/python/relax/test_codegen_tensorrt.py:67–74

(mod, inputs_np, target, legalize=False)

def build_and_run(mod, inputs_np, target, legalize=False):
    dev = tvm.device(target, 0)
    with tvm.transform.PassContext(config={"relax.transform.apply_legalize_ops": legalize}):
        ex = tvm.compile(mod, target)
    vm = relax.VirtualMachine(ex, dev)
    f = vm["main"]
    inputs = [tvm.runtime.tensor(inp, dev) for inp in inputs_np]
    return f(*inputs).numpy()


def test_tensorrt_offload():