def METHOD_NAME(device, trace=True, to_device=None):
    """test add lib with Pytorch wrapper"""
    print("\n############## Test on device:", device, "#################")
    export_dir = build_export_graph(device)
    engine = tvm.contrib.torch.GraphModule(num_inputs=2, num_outputs=1).to(device)
    x = np.random.rand(10, 5).astype("float32")
    y = np.random.rand(1, 5).astype("float32")
    expect = np.exp(y + x)
    def get_inputs_by_device(device):
        inps = [torch.Tensor(x), torch.Tensor(y)]
        if device == "cpu":
            return inps
        else:
            device_type, device_id = device.split(":")
            assert device_type == "cuda"
            return [inp.cuda(int(device_id)) for inp in inps]
    assets = [os.path.join(export_dir, i) for i in TVM_ASSETS]
    engine.init((x.shape, y.shape), *assets)
    outputs = engine.forward(get_inputs_by_device(device))
    tvm.testing.assert_allclose(outputs[0].cpu(), expect, atol=1e-5, rtol=1e-5)
    if trace:
        print("\n################ Test trace and load #################")
        scripted = torch.jit.script(engine)
        scripted_dir = tempfile.mkdtemp("scripted")
        scripted_path = os.path.join(scripted_dir, "model.pt")
        scripted.save(scripted_path)
        loaded = torch.jit.load(scripted_path)
        outputs = loaded.forward(get_inputs_by_device(device))
        tvm.testing.assert_allclose(outputs[0].cpu(), expect, atol=1e-5, rtol=1e-5)
        del scripted
        del loaded
    if to_device:
        print(
            "\n################ Test move from [{}] to [{}] #################".format(
                device, to_device
            )
        )
        engine = engine.to(to_device)
        outputs = engine.forward(get_inputs_by_device(to_device))
        tvm.testing.assert_allclose(outputs[0].cpu(), expect, atol=1e-5, rtol=1e-5)
    del engine