DiffEqFlux: SOSRA2() algorithm + BacksolveAdjoint() gives an error

Hi everyone. I have a question related to fitting a Neural SDE using DiffEqFlux.jl.

Description of the problem that is solved:
Given a specific SDE system, I want to steer that system (with a particular control) to the all-ones vector. In order to do so, I parametrize the control input as a Neural Network and try to fit the Neural Network. This NN takes as input both the position and the time. In order to train the Neural Network, I decided to use the sciml_train function of DiffEqFlux.jl in combination with the BacksolveAdjoint() method for calculating the gradient of the SDE.

Code:

using Flux, DiffEqFlux, DiffEqSensitivity, DifferentialEquations, Statistics

x_size = 6 # Size of the spatial dimensions in the SDE
v_size = 6 # Output size of the control 

# Define Neural Network for the control input
input_size = x_size + 1 # size of the spatial dimensions PLUS one time dimensions
nn_initial = Chain(Dense(input_size,v_size)) # The actual neural network
p_nn, model = Flux.destructure(nn_initial)
nn(x,p) = model(p)(x) 

# Define the right hand side of the SDE
const_mat = zeros(Float64, (x_size, v_size))
for i = 1:max(x_size,v_size)
    const_mat[i,i] = 1
end

function f!(du,u,p,t)
    du .= 2.0.*u + const_mat*nn([u;t],p)
end

function g!(du,u,p,t)
    du .= 0.0.*u .+ sqrt(2*0.001)
end

# Define SDE problem
u0 = vec(rand(Float64, (x_size,1)))
tspan = (0.0, 1.0) 
prob = SDEProblem{true}(f!, g!, u0, tspan, p_nn)

# Defining the loss function
function loss(pars, prob)

    function prob_func(prob, i, repeat)
        # Prepare new initial state and remake the problem
        u0tmp = vec(rand(Float64,(x_size,1)))

        remake(prob, p = pars, u0 = u0tmp)
    end

    ensembleprob = EnsembleProblem(prob, prob_func = prob_func)

    _sol = solve(ensembleprob, SOSRA2(), EnsembleThreads(), sensealg = BacksolveAdjoint(), saveat = vec(0:0.1:1),
    adaptive = true, trajectories = 30)

    A = convert(Array, _sol)

    return sum(abs2, A .- 1)
end

# Actually training/fitting the model
result = DiffEqFlux.sciml_train((p) -> loss(p,prob), p_nn, ADAM(0.1), maxiters=5)

Error:

ERROR: LoadError: BoundsError: attempt to access 995-element Array{Array{Float64,1},1} at index [996]
Stacktrace:
 [1] getindex at .\array.jl:809 [inlined]
 [2] #interpolate!#70 at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_process_interface.jl:524 [inlined]
 [3] #interpolate!#72 at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:15 [inlined]
 [4] NoiseWrapper at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\types.jl:179 [inlined]
 [5] calculate_step!(::NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true}, ::Float64, ::Array{Float64,1}, ::Array{Float32,1}) at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:23
 [6] setup_next_step! at C:\Users\Sven\.julia\packages\DiffEqNoiseProcess\QVesN\src\noise_interfaces\noise_wrapper_interface.jl:71 [inlined]
 [7] setup_next_step! at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\integrators\integrator_utils.jl:2 [inlined]
 [8] __init(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true},SDEFunction{true,DiffEqSensitivity.ODEBacksolveSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,ReverseDiff.GradientTape{DiffEqSensitivity.var"#89#100"{ODEFunction{true,StochasticTransformedFunction{SDEProblem{Array{Float64,1}, ............. verbose::Bool, force_dtmin::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), userdata::Nothing, initialize_integrator::Bool, seed::UInt64, alias_u0::Bool, alias_jumps::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\solve.jl:571 
 [9] #__solve#100 at C:\Users\Sven\.julia\packages\StochasticDiffEq\qg3yM\src\solve.jl:6 [inlined]
 [10] solve_call(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true},SDEFunction{true,DiffEqSensitivity.ODEBacksolveSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,ReverseDiff.GradientTape{DiffEqSensitivity.var"#89#100"{ODEFunction{true,StochasticTransformedFunction{SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1}, ......... ,Tuple{Bool,Bool,Array{Float64,1},Array{Float64,1},Float64,Float64,Bool}}}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:61
 [11] solve_up(::SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},NoiseWrapper{Float64,2,Float64,Array{Float64,1},Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},Array{Array{Float64,1},1},true},SDEFunction{true,DiffEqSensitivity.ODEBacksolveSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,ReverseDiff.GradientTape{DiffEqSensitivity.var"#89#100"{ODEFunction{true,StochasticTransformedFunction{SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(f!)t 
ypeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},SDEFunction{true,typeof(f!),typeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Array{Float64,1},Nothing},UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,typeof(SciMLBase.DEFAULT_OBSERVED),Nothing}},Tuple{ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}},ReverseDiff.TrackedArray{Float32,Float64,1,Array{Float32,1},Array{Float64,1}},ReverseDiff.TrackedArray{Float64,Float64,1,Array{Float64,1},Array{Float64,1}}}, .................................., ::SOSRA2; kwargs::Base.Iterators.Pairs{Symbol,Any,NTuple{7,Symbol},NamedTuple{(:save_everystep, :save_start, :saveat, :tstops, :abstol, :reltol, :adaptive),Tuple{Bool,Bool,Array{Float64,1},Array{Float64,1},Float64,Float64,Bool}}}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:82
 [12] #solve#57 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:70 [inlined]
 [13] _adjoint_sensitivities(::RODESolution{Float64,2,Array{Array{Float64,1},1},Nothing,Nothing,Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(f!),typeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},SOSRA2,StochasticDiffEq.LinearInterpolationData{Array{Array{Float64,1},1},Array{Float64,1}},DiffEqBase.DEStats}, ::BacksolveAdjoint{0,true,Val{:central},Bool,Bool}, ::SOSRA2, ::DiffEqSensitivity.var"#df#168"{Array{Array{Float64,1},1},Array{Float64,1},Colon}, ::Array{Float64,1}, ::Nothing; abstol::Float64, reltol::Float64, checkpoints::Array{Float64,1}, corfunc_analytical::Nothing, callback::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Bool,Tuple{Symbol},NamedTuple{(:adaptive,),Tuple{Bool}}}) at C:\Users\Sven\.julia\packages\DiffEqSensitivity\lnBE2\src\sensitivity_interface.jl:36
 [14] adjoint_sensitivities(::RODESolution{Float64,2,Array{Array{Float64,1},1},Nothing,Nothing,Array{Float64,1},NoiseProcess{Float64,2,Float64,Array{Float64,1},Array{Float64,1},Array{Array{Float64,1},1},typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_DIST),typeof(DiffEqNoiseProcess.INPLACE_WHITE_NOISE_BRIDGE),true,ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},ResettableStacks.ResettableStack{Tuple{Float64,Array{Float64,1},Array{Float64,1}},true},RSWM{Float64},Nothing,RandomNumbers.Xorshifts.Xoroshiro128Plus},SDEProblem{Array{Float64,1},Tuple{Float64,Float64},true,Array{Float32,1},Nothing,SDEFunction{true,typeof(f!),typeof(g!),UniformScaling{Bool},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},typeof(g!),Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},Nothing},SOSRA2,StochasticDiffEq.LinearInterpolationData{Array{Array{Float64,1},1},Array{Float64,1}},DiffEqBase.DEStats}, ::SOSRA2, ::Vararg{Any,N} where N; sensealg::BacksolveAdjoint{0,true,Val{:central},Bool,Bool}, kwargs::Base.Iterators.Pairs{Symbol,Union{Nothing, Bool},Tuple{Symbol,Symbol},NamedTuple{(:callback, :adaptive),Tuple{Nothing,Bool}}}) at C:\Users\Sven\.julia\packages\DiffEqSensitivity\lnBE2\src\sensitivity_interface.jl:6
 [15] (::DiffEqSensitivity.var"#adjoint_sensitivity_backpass#167"{Base.Iterators.Pairs{Symbol,Bool,Tuple{Symbol},NamedTuple{(:adaptive,),Tuple{Bool}}},SOSRA2,BacksolveAdjoint{0,true,Val{:central},Bool,Bool},Array{Float64,1},Array{Float32,1},Tuple{},NamedTuple{(:adaptive,),Tuple{Bool}},Colon})(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\DiffEqSensitivity\lnBE2\src\concrete_solve.jl:179
 [16] #263#back at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:65 [inlined]
 [17] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
 [18] (::Zygote.var"#1701#back#182"{Zygote.var"#180#181"{DiffEqBase.var"#263#back#74"{DiffEqSensitivity.var"#adjoint_sensitivity_backpass#167"{Base.Iterators.Pairs{Symbol,Bool,Tuple{Symbol},NamedTuple{(:adaptive,),Tuple{Bool}}},SOSRA2,BacksolveAdjoint{0,true,Val{:central},Bool,Bool},Array{Float64,1},Array{Float32,1},Tuple{},NamedTuple{(:adaptive,),Tuple{Bool}},Colon}},Tuple{NTuple{6,Nothing},Tuple{Nothing}}}})(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59
 [19] #solve#57 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:70 [inlined]
 [20] (::typeof(∂(#solve#57)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [21] (::Zygote.var"#180#181"{typeof(∂(#solve#57)),Tuple{NTuple{6,Nothing},Tuple{Nothing}}})(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194
 [22] (::Zygote.var"#1701#back#182"{Zygote.var"#180#181"{typeof(∂(#solve#57)),Tuple{NTuple{6,Nothing},Tuple{Nothing}}}})(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59
 [23] (::typeof(∂(solve##kw)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [24] #batch_func#452 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:143 [inlined]
 [25] (::typeof(∂(#batch_func#452)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
 [26] #457 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:195 [inlined]
 [27] (::typeof(∂(λ)))(::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [28] (::DiffEqBase.var"#219#227")(::typeof(∂(λ)), ::Array{Array{Float64,1},1}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\init.jl:259
 [29] responsible_map(::Function, ::Array{typeof(∂(λ)),1}, ::Vararg{Any,N} where N) at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:188
 [30] (::DiffEqBase.var"#∇responsible_map_internal#226"{Array{typeof(∂(λ)),1}})(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\init.jl:259
 [31] #361#back at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59 [inlined]
 [32] #solve_batch#456 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:194 [inlined]
 [33] (::typeof(∂(#solve_batch#456)))(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
 [34] #solve_batch#459 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:203 [inlined]
 [35] (::typeof(∂(#solve_batch#459)))(::EnsembleSolution{Array{Float64,1},2,Array{Array{Array{Float64,1},1},1}}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
 [36] macro expansion at .\timing.jl:233 [inlined]
 [37] #__solve#451 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\ensemble\basic_ensemble_solve.jl:108 [inlined]
 [38] (::typeof(∂(#__solve#451)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0 (repeats 2 times)
 [39] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
 [40] (::Zygote.var"#1701#back#182"{Zygote.var"#180#181"{typeof(∂(__solve##kw)),Tuple{Tuple{Nothing,Nothing,Nothing},Tuple{Nothing,Nothing}}}})(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59
 [41] #solve#59 at C:\Users\Sven\.julia\packages\DiffEqBase\U3Zj7\src\solve.jl:96 [inlined]
 [42] (::typeof(∂(#solve#59)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [43] (::Zygote.var"#180#181"{typeof(∂(#solve#59)),Tuple{Tuple{Nothing,Nothing,Nothing},Tuple{Nothing,Nothing}}})(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194
 [44] #1701#back at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59 [inlined]
 [45] (::typeof(∂(solve##kw)))(::Array{Float64,3}) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [46] loss at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:43 [inlined]
 [47] (::typeof(∂(loss)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [48] #39 at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:52 [inlined]
 [49] (::typeof(∂(#39)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [50] #69 at C:\Users\Sven\.julia\packages\DiffEqFlux\alPQ3\src\train.jl:3 [inlined]
 [51] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
 [52] #1701#back at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59 [inlined]
 [53] OptimizationFunction at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\problems\basic_problems.jl:107 [inlined]
 ... (the last 3 lines are repeated 1 more time)
 [57] #180 at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\lib\lib.jl:194 [inlined]
 [58] (::Zygote.var"#1701#back#182"{Zygote.var"#180#181"{typeof(∂(λ)),Tuple{Tuple{Nothing,Nothing},Int64}}})(::Float64) at C:\Users\Sven\.julia\packages\ZygoteRules\OjfTt\src\adjoint.jl:59
 [59] #8 at C:\Users\Sven\.julia\packages\GalacticOptim\JnLwV\src\solve.jl:94 [inlined]
 [60] (::typeof(∂(λ)))(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface2.jl:0
 [61] (::Zygote.var"#69#70"{Zygote.Params,Zygote.Context,typeof(∂(λ))})(::Float64) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface.jl:252
 [62] gradient(::Function, ::Zygote.Params) at C:\Users\Sven\.julia\packages\Zygote\6HN9x\src\compiler\interface.jl:59
 [63] __solve(::OptimizationProblem{false,OptimizationFunction{false,GalacticOptim.AutoZygote,OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},GalacticOptim.var"#146#156"{GalacticOptim.var"#145#155"{OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Nothing}},GalacticOptim.var"#149#159"{GalacticOptim.var"#145#155"{OptimizationFunction{true,GalacticOptim.AutoZygote,DiffEqFlux.var"#69#70"{var"#39#40"},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Nothing}},GalacticOptim.var"#154#164",Nothing,Nothing,Nothing},Array{Float32,1},SciMLBase.NullParameters,Nothing,Nothing,Nothing,Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxiters,),Tuple{Int64}}}}, ::ADAM, ::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; maxiters::Int64, cb::Function, progress::Bool, save_best::Bool, kwargs::Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}) at C:\Users\Sven\.julia\packages\GalacticOptim\JnLwV\src\solve.jl:93
 [64] #solve#468 at C:\Users\Sven\.julia\packages\SciMLBase\XuLdB\src\solve.jl:3 [inlined]
 [65] sciml_train(::var"#39#40", ::Array{Float32,1}, ::ADAM, ::GalacticOptim.AutoZygote; lower_bounds::Nothing, upper_bounds::Nothing, kwargs::Base.Iterators.Pairs{Symbol,Int64,Tuple{Symbol},NamedTuple{(:maxiters,),Tuple{Int64}}}) at C:\Users\Sven\.julia\packages\DiffEqFlux\alPQ3\src\train.jl:6
in expression starting at c:\Users\Sven\Documents\Master\Master_thesis\Julia\Code\TestRepository\TestFileDiscourse_question2.jl:52

Small remark
I removed a part of the error at lines [8], [10] and [11]. The reason is that it was very large and made me go over the character limit for the post. The parts that I removed are replaced by dots (so ‘…’).

Possible location of the error
At the top line of the error, I see ERROR: LoadError: BoundsError: attempt to access 995-element Array{Array{Float64,1},1} at index [996]. Here the values 995 and 996 change from run to run. Hence, most likely the error is related to the stochasticity and adaptivity when solving a SDE (either the forward SDE or the Backward SDE of BacksolveAdjoint()). Moreover, when changing the SOSRA2() algorithm to the LambaEulerHeun() algorithm, no errors occur. Hence, I suspect that the error is related to the SOSRA2() algorithm and related algorithms (for instance, with SRA2() the code gives the same error).

Question:
Is there any way I can change the code such that the SOSRA2() algorithm can be used in combination with BacksolveAdjoint()?

Fixed by https://github.com/SciML/DiffEqNoiseProcess.jl/pull/93

Thank you very much for your reply. After updating the package, I used SOSRA2() on some newly made code and it worked.