From f6a0aa2a5535ada3d9008eaecd8db787942dcec2 Mon Sep 17 00:00:00 2001 From: optima2005 <56945758+optima2005@users.noreply.github.com> Date: Fri, 25 Oct 2019 00:56:00 +0800 Subject: [PATCH] [Documentation]Fix example code in comment of tvm.build_module.build() (#4195) * Fix example code in comment of tvm.build_module.build() * Update build_module.py --- python/tvm/build_module.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py index 4cb0993..479bdce 100644 --- a/python/tvm/build_module.py +++ b/python/tvm/build_module.py @@ -568,10 +568,11 @@ def build(inputs, B = tvm.placeholder((n,), name='B') C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') s1 = tvm.create_schedule(C.op) - s2 = topi.cpp.cuda.schedule_injective("cuda", [C]) - f1 = tvm.lower(s1, [A, B, C], name="test_add1") - f2 = tvm.lower(s2, [A, B, C], name="test_add2") - m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm") + with tvm.target.cuda() as cuda_tgt: + s2 = topi.cuda.schedule_injective(cuda_tgt, [C]) + f1 = tvm.lower(s1, [A, B, C], name="test_add1") + f2 = tvm.lower(s2, [A, B, C], name="test_add2") + m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm") Note ---- -- 2.7.4