fix: fix observed variable adjoint

SciML · Oct 7, 2024 · 6184c69 · 6184c69
1 parent 9b96078
commit 6184c69
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/ext/SciMLBaseZygoteExt.jl b/ext/SciMLBaseZygoteExt.jl
@@ -106,11 +106,12 @@ end
         if is_observed(VA, sym)
             f = observed(VA, sym)
             p = parameter_values(VA)
-            tunables, _, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
+            tunables, repack, _ = SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)
             u = state_values(VA)
             t = current_time(VA)
             y, back = Zygote.pullback(u, tunables) do u, tunables
-                f.(u, Ref(tunables), t)
+                _p = repack(tunables)
+                f.(u, Ref(_p), t)
             end
             gs = back(Δ)
             (gs[1], nothing)