Skip to content

Commit a52552c

Browse files
fix: handle rrule returning extra value in pullback
1 parent d53a127 commit a52552c

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

ext/SciMLBaseZygoteExt.jl

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,8 @@ end
4444
struct ZygoteConfig <: ChainRulesCore.RuleConfig{ChainRulesCore.HasReverseMode} end
4545

4646
@adjoint function Base.getindex(VA::ODESolution, sym, j::Integer)
47-
Zygote.ChainRulesCore.rrule(ZygoteConfig(), getindex, VA, sym, j)
47+
res, pullback = Zygote.ChainRulesCore.rrule(ZygoteConfig(), getindex, VA, sym, j)
48+
return res, Base.tail pullback
4849
end
4950

5051
@adjoint function EnsembleSolution(sim, time, converged, stats)

0 commit comments

Comments
 (0)