diff --git a/python/tvm/build_module.py b/python/tvm/build_module.py
index 4cb09931616e..479bdcead7c5 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
     ----