Allow for Enzyme choice with callbacks automatically

SciML · Jun 7, 2024 · 33a57d2 · 33a57d2
1 parent 2a110dc
commit 33a57d2
Showing 1 changed file with 4 additions and 7 deletions.
diff --git a/src/sensitivity_interface.jl b/src/sensitivity_interface.jl
@@ -364,14 +364,11 @@ function adjoint_sensitivities(sol, args...;
         else
             has_cb = false
         end
-        if !has_cb
-            _sensealg = if isinplace(sol.prob)
-                setvjp(sensealg, inplace_vjp(sol.prob, sol.prob.u0, sol.prob.p, verbose))
-            else
-                setvjp(sensealg, ZygoteVJP())
-            end
+
+        _sensealg = if isinplace(sol.prob)
+          setvjp(sensealg, inplace_vjp(sol.prob, sol.prob.u0, sol.prob.p, verbose))
         else
-            _sensealg = setvjp(sensealg, ReverseDiffVJP())
+            setvjp(sensealg, ZygoteVJP())
         end
 
         return try