diff --git a/tests/python/unittest/test_tir_ir_builder.py b/tests/python/unittest/test_tir_ir_builder.py index e2a277f9f307b..df8b6e9979163 100644 --- a/tests/python/unittest/test_tir_ir_builder.py +++ b/tests/python/unittest/test_tir_ir_builder.py @@ -511,7 +511,7 @@ def test_device_ir(A, B): tx = te.thread_axis("threadIdx.x") ib.scope_attr(tx, "thread_extent", n) - temp = ib.allocate(dtype, (n,), scope="dyn.shared") + temp = ib.allocate(dtype, (n,), scope="dyn.shared") # n is symbolic size Aptr = ib.buffer_ptr(A) Bptr = ib.buffer_ptr(B) @@ -538,8 +538,8 @@ def test_device_ir(A, B): s = te.create_schedule(B.op) def check_target(target): - # if not tvm.testing.device_enabled(target): - # return + if not tvm.testing.device_enabled(target): + return freduce = tvm.build(s, [A, B], target) dev = tvm.device(target, 0)