Commit f6a0aa2a by optima2005 Committed by Tianqi Chen

[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
parent c9aa55cd
...@@ -568,7 +568,8 @@ def build(inputs, ...@@ -568,7 +568,8 @@ def build(inputs,
B = tvm.placeholder((n,), name='B') B = tvm.placeholder((n,), name='B')
C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C') C = tvm.compute(A.shape, lambda *i: A(*i) + B(*i), name='C')
s1 = tvm.create_schedule(C.op) s1 = tvm.create_schedule(C.op)
s2 = topi.cpp.cuda.schedule_injective("cuda", [C]) 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") f1 = tvm.lower(s1, [A, B, C], name="test_add1")
f2 = tvm.lower(s2, [A, B, C], name="test_add2") f2 = tvm.lower(s2, [A, B, C], name="test_add2")
m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm") m = tvm.build({"llvm": [f1], "cuda": [f2]}, target_host="llvm")
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment