DiffEqFlux.sciml_train fails with save_idxs keyword

Training via DiffEqFlux fails when save_idxs keyword is used within the loss function. Differentiation seems to be the problem. Evaluation of the loss function works fine.

MWE is taken from the docs

using DifferentialEquations, Flux, Optim, DiffEqFlux

function lotka_volterra!(du, u, p, t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = α*x - β*x*y
  du[2] = dy = -δ*y + γ*x*y
end

# Initial condition
u0 = [1.0, 1.0]

# Simulation interval and intermediary points
tspan = (0.0, 10.0)
tsteps = 0.0:0.1:10.0

# LV equation parameter. p = [α, β, δ, γ]
p = [1.5, 1.0, 3.0, 1.0]

# Setup the ODE problem, then solve
prob = ODEProblem(lotka_volterra!, u0, tspan, p)

function loss(p)
  sol = solve(prob, Tsit5(), p=p, save_idxs=[2], saveat = tsteps)
  loss = sum(abs2, sol.-1)
  return loss, sol
end


result_ode = DiffEqFlux.sciml_train(loss, p, ADAM(0.1), maxiters = 100)

Open an issue.

Done, see here.

Cool a fix is in Fix array save_idxs by ChrisRackauckas · Pull Request #408 · SciML/DiffEqSensitivity.jl · GitHub