Hey, I’m trying to implement minibatching in a Neural ODE system. I calculate the gradient
for each batch with the following code:

``````
using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots, LinearAlgebra, Distributions

ode_data = Array(solve(prob_trueode, Tsit5(), saveat = tsteps))

dudt2 = FastChain(FastDense(2, 50, tanh), FastDense(50, 2))
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)
train_node = NeuralODE(dudt2, (0., 1.), Tsit5(), saveat = tsteps[1:40])
y_train = ode_data[:, 1:250]
training_span = tsteps[1:250]

function predict(p, time_batch)
prob = NeuralODE(dudt2, (tsteps[1], tsteps[end]), Tsit5(), saveat = time_batch)
prob(u0, p)
end

function loss(p, batch, time_batch)
pred = predict(p, time_batch)
sum(abs2, pred .- batch)
end

for i in 1:length(t_batch)
end
end

``````

I need to calculate the gradient of this function with respect to θ, but when I do:

``````
ERROR: BoundsError
Stacktrace:
[1] macro expansion at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0 [inlined]
[2] _pullback(::Zygote.Context, ::typeof(throw), ::BoundsError) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:12
[3] getindex at ./number.jl:78 [inlined]
[4] literal_getindex at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/tools/builtins.jl:15 [inlined]
[5] _pullback(::Zygote.Context, ::typeof(Zygote.literal_getindex), ::Float32, ::Val{2}) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[6] loss at ./REPL[15]:3 [inlined]
[7] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::typeof(loss), ::Array{Float32,1}, ::Array{Float32,1}, ::Array{Float64,1}) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[8] #11 at ./REPL[23]:4 [inlined]
[9] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::Zygote.Context, ::var"#11#12"{Array{Float32,2},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Int64}, ::Array{Float32,1}) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[12] _pullback at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:38 [inlined]
[13] _pullback(::Zygote.Context, ::typeof(ZygoteRules._pullback), ::var"#11#12"{Array{Float32,2},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Int64}, ::Array{Float32,1}) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[16] pullback at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:44 [inlined]
[17] _pullback(::Zygote.Context, ::typeof(ZygoteRules.pullback), ::var"#11#12"{Array{Float32,2},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},Int64}, ::Array{Float32,1}) at /home/aslan_garcia/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0