remove syncing at each eval point

martenlienen · Oct 2, 2024 · 551942f · 551942f
1 parent d50855c
commit 551942f
Showing 1 changed file with 9 additions and 10 deletions.
diff --git a/torchode/adjoints.py b/torchode/adjoints.py
@@ -220,17 +220,16 @@ def solve(
                     )
                     >= 0.0
                 ) & not_yet_evaluated
-                if to_be_evaluated.any():
-                    interpolation = step_method.build_interpolation(interp_data)
-                    nonzero = to_be_evaluated.nonzero()
-                    sample_idx, eval_t_idx = nonzero[:, 0], nonzero[:, 1]
-                    y_eval[sample_idx, eval_t_idx] = interpolation.evaluate(
-                        t_eval[sample_idx, eval_t_idx], sample_idx
-                    )
+                interpolation = step_method.build_interpolation(interp_data)
+                nonzero = to_be_evaluated.nonzero()
+                sample_idx, eval_t_idx = nonzero[:, 0], nonzero[:, 1]
+                y_eval[sample_idx, eval_t_idx] = interpolation.evaluate(
+                    t_eval[sample_idx, eval_t_idx], sample_idx
+                )
 
-                    not_yet_evaluated = torch.logical_xor(
-                        to_be_evaluated, not_yet_evaluated
-                    )
+                not_yet_evaluated = torch.logical_xor(
+                    to_be_evaluated, not_yet_evaluated
+                )
 
             ########################
             # Update the step size #