def METHOD_NAME():
    """main test function"""
    def build_export_vm(device):
        """relay build & export graph"""
        x = relay.var("x", shape=(10, 5))
        y = relay.var("y", shape=(1, 5))
        z = relay.add(x, y)
        z = relay.exp(z)
        func = relay.Function([x, y], z)
        x_data = np.random.rand(10, 5).astype("float32")
        y_data = np.random.rand(1, 5).astype("float32")
        pt_device = torch.device(device)
        if pt_device.type == "cuda":
            target = "cuda"
            ctx = tvm.cuda(pt_device.index)
        else:
            target = "llvm"
            ctx = tvm.cpu(0)
        exe = relay.vm.compile(tvm.IRModule.from_expr(func), target=target, params={})
        code, lib = exe.save()
        export_dir = tempfile.mkdtemp("tvm_export")
        # export to tempdir
        lib.export_library(os.path.join(export_dir, TVM_ASSETS[0]))
        with open(os.path.join(export_dir, TVM_ASSETS[1]), "wb") as fout:
            fout.write(code)
        vm = tvm.runtime.vm.VirtualMachine(exe, ctx)
        res = vm.run(x_data, y_data)
        ref_res = np.exp(y_data + x_data)
        tvm.testing.assert_allclose(res.numpy(), ref_res, atol=1e-5, rtol=1e-5)
        return export_dir
    def test_pt_run(device, trace=True, to_device=None, inp_on_cuda=False):
        """test add lib with Pytorch wrapper"""
        print("\n############## Test on device:", device, "#################")
        export_dir = build_export_vm(device)
        engine = tvm.contrib.torch.VMModule(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
    test_pt_run(device="cuda:0", trace=True, to_device="cuda:1", inp_on_cuda=True)
    test_pt_run(device="cpu", trace=True, inp_on_cuda=False)