update

loganbvh · Sep 18, 2023 · 4a4c9de · 4a4c9de
1 parent 0832ca2
commit 4a4c9de
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 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 cupyx.jit.range(2):
+        for k in range(2):
             tmp = 0.0
-            for j in cupyx.jit.range(sites.shape[0]):
+            for j in 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)

diff --git a/tdgl/solver/solve.py b/tdgl/solver/solve.py
@@ -369,7 +369,6 @@ def update(
                 get_A_induced_cupy[128, 1024](
                     J_site, areas, sites, edge_centers, new_A_induced
                 )
-                cupy.cuda.Stream.null.synchronize()
             else:
                 new_A_induced = get_A_induced_numba(J_site, areas, sites, edge_centers)
             # Update induced vector potential using Polyak's method
@@ -381,6 +380,7 @@ def update(
                 screening_error = xp.max(
                     xp.linalg.norm(dA, axis=1) / xp.linalg.norm(A_induced, axis=1)
                 )
+                screening_error = float(screening_error)
                 del v[:-2]
                 del A_induced_vals[:-2]