Update screening.py

loganbvh · Sep 18, 2023 · c63bf48 · c63bf48
1 parent 1cc0b29
commit c63bf48
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tdgl/solver/screening.py b/tdgl/solver/screening.py
@@ -59,9 +59,9 @@ def get_A_induced_cupy(
         A_induced,
     ):
         i = cupyx.jit.grid(1)
-        for k in range(2):
+        for k in cupyx.jit.range(2):
             tmp = 0.0
-            for j in range(sites.shape[0]):
+            for j in cupyx.jit.range(sites.shape[0]):
                 dx = edge_centers[i, 0] - sites[j, 0]
                 dy = edge_centers[i, 1] - sites[j, 1]
                 dr = cupy.sqrt(dx * dx + dy * dy)