Skip to content

Commit

Permalink
[Documentation]Fix example code in comment of tvm.build_module.build() (
Browse files Browse the repository at this point in the history
apache#4195)

* Fix example code in comment of tvm.build_module.build()

* Update build_module.py
  • Loading branch information
optima2005 authored and kevinthesun committed Oct 30, 2019
1 parent 15ff567 commit 07e3a1b
Showing 1 changed file with 5 additions and 4 deletions.
9 changes: 5 additions & 4 deletions python/tvm/build_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
----
Expand Down

0 comments on commit 07e3a1b

Please sign in to comment.