Issue with custom Flux loss function

General Usage
1

Hi I am trying to create a custom loss function in Flux which uses the output from the flux model as parameters to an ODE equation. However I am getting an error that I have been unable to solve.

The code is the following

using Flux
using DifferentialEquations
using SciMLSensitivity
using RecursiveArrayTools
using Zygote

function DynamicAeroXvar!(qdot, q, p,t)
    ~, Boosting = p
    qdot=5*q.+Boosting;
    return qdot
end

CorrectShapeMat=rand(1,20)*5
train_yGradBoost=rand(1,20)*10
model = Chain(Dense(1, 4, elu),Dense(4, 1))
data = Flux.DataLoader((CorrectShapeMat, train_yGradBoost), batchsize=4, shuffle=true)
pars = Flux.params(model)  # contains references to arrays in model
opt = Flux.Adam(0.01)      # will store optimiser momentum, etc.

# Training loop, using the whole data set 1000 times:
#evalcb() = @show(loss(test_x, test_y))
#throttled_cb = throttle(evalcb, 5)
for epoch in 1:10
    Flux.train!(pars, data, opt) do x, y
        
    ~, batchSize = size(y)         
    NeuralNetPrediction = model(x) 
    outCoef=zeros(1,batchSize)
    outCoefBuffed = Zygote.Buffer(outCoef)
        
    for i=1:batchSize
        GradBoost = NeuralNetPrediction[i] .*1.0
        a=(2.0, GradBoost)
        p = ArrayPartition(a)
            
        tspan = (0.0, 0.001)
        prob = ODEProblem(DynamicAeroXvar!,[1.0],tspan,p)
        
        sol = solve(prob)
        X0=sol[1,end]
        outCoefBuffed[1, i] = X0 .*100
        #push!(Pusharray,[LiftCoefs1])
    end

    NeuralCoef = copy(outCoefBuffed)
    Flux.mse(NeuralCoef, y)  
    end
end

This produces the following error message.

ERROR: MethodError: no method matching similar(::Float64, ::Tuple{Base.OneTo{Int64}})
Closest candidates are:
  similar(::Type{<:StructArrays.StructArray{T, N, C}}, ::Tuple{Union{Integer, Base.OneTo}, Vararg{Union{Integer, Base.OneTo}}}) where {T, N, C} at C:\Users\smickus3\.julia\packages\StructArrays\wkp4X\src\structarray.jl:293
  similar(::Type{<:StructArrays.StructArray{T, N, C}}, ::Tuple{Union{Integer, AbstractUnitRange}, Vararg{Union{Integer, AbstractUnitRange}}}) where {T, N, C} at C:\Users\smickus3\.julia\packages\StructArrays\wkp4X\src\structarray.jl:293  similar(::Type{T}, ::Tuple{Union{Integer, Base.OneTo}, Vararg{Union{Integer, Base.OneTo}}}) where T<:AbstractArray at abstractarray.jl:840
  ...
Stacktrace:
  [1] _unsafe_getindex(#unused#::IndexCartesian, A::ArrayPartition{Float64, Tuple{Float64, Float64}}, I::UnitRange{Int64})
    @ RecursiveArrayTools C:\Users\smickus3\.julia\packages\RecursiveArrayTools\gH9Kb\src\array_partition.jl:244      
  [2] _getindex
    @ .\multidimensional.jl:875 [inlined]
  [3] getindex
    @ .\abstractarray.jl:1241 [inlined]
  [4] (::SciMLSensitivity.var"#290#299"{0, VectorOfArray{Float64, 2, Vector{Vector{Float64}}}, Vector{Float64}, Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:verbose,), Tuple{Bool}}}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ArrayPartition{Float64, Tuple{Float64, Float64}}, ODEFunction{true, SciMLBase.AutoSpecialize, FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{Float64}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}}, false}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Nothing, ForwardDiffSensitivity{0, nothing}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Tuple{}, Vector{Float64}})()
    @ SciMLSensitivity C:\Users\smickus3\.julia\packages\SciMLSensitivity\DInxI\src\concrete_solve.jl:651
  [5] unthunk
    @ C:\Users\smickus3\.julia\packages\ChainRulesCore\C73ay\src\tangent_types\thunks.jl:204 [inlined]
  [6] wrap_chainrules_output
    @ C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\chainrules.jl:105 [inlined]
  [7] map
    @ .\tuple.jl:223 [inlined]
  [8] map (repeats 3 times)
    @ .\tuple.jl:224 [inlined]
  [9] wrap_chainrules_output
    @ C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\chainrules.jl:106 [inlined]
 [10] ZBack
    @ C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\chainrules.jl:206 [inlined]
 [11] (::Zygote.var"#208#209"{Tuple{NTuple{4, Nothing}, Tuple{}}, Zygote.ZBack{SciMLSensitivity.var"#forward_sensitivity_backpass#298"{0, Vector{Float64}, Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:verbose,), Tuple{Bool}}}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ArrayPartition{Float64, Tuple{Float64, Float64}}, ODEFunction{true, SciMLBase.AutoSpecialize, FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{Float64}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}}, false}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, 
Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, 
Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Nothing, ForwardDiffSensitivity{0, nothing}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, SciMLBase.ChainRulesOriginator, Tuple{}, Vector{Float64}}}})(Δ::VectorOfArray{Float64, 2, Vector{Vector{Float64}}})
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\lib\lib.jl:206
 [12] (::Zygote.var"#2066#back#210"{Zygote.var"#208#209"{Tuple{NTuple{4, Nothing}, Tuple{}}, Zygote.ZBack{SciMLSensitivity.var"#forward_sensitivity_backpass#298"{0, Vector{Float64}, Base.Pairs{Symbol, Bool, Tuple{Symbol}, NamedTuple{(:verbose,), Tuple{Bool}}}, ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, true, ArrayPartition{Float64, Tuple{Float64, Float64}}, ODEFunction{true, SciMLBase.AutoSpecialize, FunctionWrappersWrappers.FunctionWrappersWrapper{Tuple{FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{Float64}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, 
Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, Float64}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}, FunctionWrappers.FunctionWrapper{Nothing, Tuple{Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, Vector{ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}, ArrayPartition{Float64, Tuple{Float64, Float64}}, ForwardDiff.Dual{ForwardDiff.Tag{DiffEqBase.OrdinaryDiffEqTag, Float64}, Float64, 1}}}}, false}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, 
Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing}, Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}, SciMLBase.StandardODEProblem}, Nothing, ForwardDiffSensitivity{0, nothing}, Vector{Float64}, ArrayPartition{Float64, Tuple{Float64, Float64}}, SciMLBase.ChainRulesOriginator, Tuple{}, Vector{Float64}}}}})(Δ::VectorOfArray{Float64, 2, Vector{Vector{Float64}}})     
    @ Zygote C:\Users\smickus3\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67
 [13] Pullback
    @ C:\Users\smickus3\.julia\packages\DiffEqBase\WXn2i\src\solve.jl:802 [inlined]
 [14] (::typeof(∂(#solve#31)))(Δ::VectorOfArray{Float64, 2, Vector{Vector{Float64}}})
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface2.jl:0
 [15] #208
    @ C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\lib\lib.jl:206 [inlined]
 [16] #2066#back
    @ C:\Users\smickus3\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67 [inlined]
 [17] Pullback
    @ C:\Users\smickus3\.julia\packages\DiffEqBase\WXn2i\src\solve.jl:792 [inlined]
 [18] (::typeof(∂(solve)))(Δ::VectorOfArray{Float64, 2, Vector{Vector{Float64}}})
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface2.jl:0
 [19] Pullback
    @ c:\Users\smickus3\OneDrive - Georgia Institute of Technology\CAMM_Lab_Research\ONR_Bleed_Air\GradientBoostingCoefficientDetermination.jl:39 [inlined]
 [20] (::typeof(∂(#31)))(Δ::Float64)
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface2.jl:0
 [21] #208
    @ C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\lib\lib.jl:206 [inlined]
 [22] #2066#back
    @ C:\Users\smickus3\.julia\packages\ZygoteRules\AIbCs\src\adjoint.jl:67 [inlined]
 [23] Pullback
    @ C:\Users\smickus3\.julia\packages\Flux\FKl3M\src\optimise\train.jl:132 [inlined]
 [24] (::typeof(∂(λ)))(Δ::Float64)
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface2.jl:0
 [25] (::Zygote.var"#99#100"{Params{Zygote.Buffer{Any, Vector{Any}}}, typeof(∂(λ)), Zygote.Context{true}})(Δ::Float64)    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface.jl:389
 [26] withgradient(f::Function, args::Params{Zygote.Buffer{Any, Vector{Any}}})
    @ Zygote C:\Users\smickus3\.julia\packages\Zygote\SmJK6\src\compiler\interface.jl:133
 [27] macro expansion
    @ C:\Users\smickus3\.julia\packages\Flux\FKl3M\src\optimise\train.jl:131 [inlined]
 [28] macro expansion
    @ C:\Users\smickus3\.julia\packages\ProgressLogging\6KXlp\src\ProgressLogging.jl:328 [inlined]
 [29] train!(loss::Function, ps::Params{Zygote.Buffer{Any, Vector{Any}}}, data::MLUtils.DataLoader{Tuple{Matrix{Float64}, Matrix{Float64}}, Random._GLOBAL_RNG, Val{nothing}}, opt::Adam; cb::Flux.Optimise.var"#38#41")
    @ Flux.Optimise C:\Users\smickus3\.julia\packages\Flux\FKl3M\src\optimise\train.jl:129
 [30] train!(loss::Function, ps::Params{Zygote.Buffer{Any, Vector{Any}}}, data::MLUtils.DataLoader{Tuple{Matrix{Float64}, Matrix{Float64}}, Random._GLOBAL_RNG, Val{nothing}}, opt::Adam)
    @ Flux.Optimise C:\Users\smickus3\.julia\packages\Flux\FKl3M\src\optimise\train.jl:125
 [31] top-level scope
    @ c:\Users\smickus3\OneDrive - Georgia Institute of Technology\CAMM_Lab_Research\ONR_Bleed_Air\GradientBoostingCoefficientDetermination.jl:24

Any help would be greatly appreciated.