Update driver.py

IntelPython · Nov 2, 2023 · 5777b67 · 5777b67
1 parent 9b03222
commit 5777b67
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/driver.py b/driver.py
@@ -1,8 +1,8 @@
 import dpctl
 import dpnp
-from numba import njit
 
-from numba_dpex import dpjit, kernel
+import numba_dpex.experimental as nd_exp
+from numba_dpex import Range, dpjit
 from numba_dpex.core.kernel_interface.dpcpp_iface import (
     AddressSpace,
     AtomicRef,
@@ -12,7 +12,7 @@
 from numba_dpex.utils.constants import address_space
 
 
-@dpjit
+@nd_exp.kernel
 def test_atomic_ref(a, b):
     # i = get_global_id(0)
     v = AtomicRef(
@@ -28,5 +28,5 @@ def test_atomic_ref(a, b):
 a = dpnp.ones(10, sycl_queue=q, dtype=dpnp.int64)
 b = dpnp.array(0, sycl_queue=q, dtype=dpnp.int64)
 
-test_atomic_ref(a, b)
+nd_exp.call_kernel(test_atomic_ref, Range(10), a, b)
 print(b)