| 98 | @pytest.mark.gpu |
| 99 | @pytest.mark.skipif(not env.has_cuda(), reason="need cuda") |
| 100 | def test_vm_run(): |
| 101 | mod = Module |
| 102 | target = tvm.target.Target("cuda", host="llvm") |
| 103 | ex = codegen(mod, target) |
| 104 | dev = tvm.cuda(0) |
| 105 | vm = relax.VirtualMachine(ex, dev) |
| 106 | x_np = np.random.uniform(size=(16, 16)).astype("float32") |
| 107 | x = tvm.runtime.tensor(x_np, dev) |
| 108 | y = vm["main"](x) |
| 109 | y_np = x_np + 1.0 + 1.0 + 1.0 + 1.0 |
| 110 | tvm.testing.assert_allclose(y.numpy(), y_np, rtol=1e-5, atol=1e-5) |
| 111 | |
| 112 | |
| 113 | @pytest.mark.gpu |