Update solver.py

loganbvh · Sep 22, 2023 · 34dae37 · 34dae37
1 parent 74d20d3
commit 34dae37
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/tdgl/solver/solver.py b/tdgl/solver/solver.py
@@ -447,17 +447,32 @@ def update(
                 psi_laplacian,
                 options,
             )
+
+            if use_cupy:
+                cupy.cuda.get_current_stream().synchronize()
+
             # Compute the supercurrent, scalar potential, and normal current
             supercurrent = operators.get_supercurrent(psi)
+
+            if use_cupy:
+                cupy.cuda.get_current_stream().synchronize()
+
             rhs = (divergence @ (supercurrent - dA_dt)) - (
                 mu_boundary_laplacian @ mu_boundary
             )
             if options.sparse_solver is SparseSolver.PARDISO:
                 mu = pypardiso.spsolve(mu_laplacian, rhs)
             else:
                 mu = mu_laplacian_lu(rhs)
+
+            if use_cupy:
+                cupy.cuda.get_current_stream().synchronize()
+
             normal_current = -(mu_gradient @ mu) - dA_dt
 
+            if use_cupy:
+                cupy.cuda.get_current_stream().synchronize()
+
             if not options.include_screening:
                 break