From 33a57d2b3c6ef846bd536bb9452b7667d6bedcac Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 29 May 2024 08:58:25 +0200 Subject: [PATCH] Allow for Enzyme choice with callbacks automatically --- src/sensitivity_interface.jl | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/src/sensitivity_interface.jl b/src/sensitivity_interface.jl index b5fa3e45c..b438e11e9 100644 --- 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