Skip to content

Commit 07e3a1b

Browse files
optima2005kevinthesun
authored andcommitted
[Documentation]Fix example code in comment of tvm.build_module.build() (apache#4195)
* Fix example code in comment of tvm.build_module.build() * Update build_module.py
1 parent 15ff567 commit 07e3a1b

File tree

1 file changed

+5
-4
lines changed

1 file changed

+5
-4
lines changed

python/tvm/build_module.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -568,10 +568,11 @@ def build(inputs,
568568
B = tvm.placeholder((n,), name='B')
569569
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
570570
s1 = tvm.create_schedule(C.op)
571-
s2 = topi.cpp.cuda.schedule_injective("cuda", [C])
572-
f1 = tvm.lower(s1, [A, B, C], name="test_add1")
573-
f2 = tvm.lower(s2, [A, B, C], name="test_add2")
574-
m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm")
571+
with tvm.target.cuda() as cuda_tgt:
572+
s2 = topi.cuda.schedule_injective(cuda_tgt, [C])
573+
f1 = tvm.lower(s1, [A, B, C], name="test_add1")
574+
f2 = tvm.lower(s2, [A, B, C], name="test_add2")
575+
m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm")
575576
576577
Note
577578
----

0 commit comments

Comments
 (0)