NEURAL ODE example with GPU not working

[matutenun]

the following example from https://diffeqflux.sciml.ai/dev/GPUs/ is not working.

Here is the full neural ODE example. Note that we use the gpu function so that the same code works on CPUs and GPUs, dependent on using CUDA.

using DiffEqFlux, OrdinaryDiffEq, Flux, Optim, Plots, CUDA, DiffEqSensitivity
CUDA.allowscalar(false) # Makes sure no slow operations are occuring

# Generate Data
u0 = Float32[2.0; 0.0]
datasize = 30
tspan = (0.0f0, 1.5f0)
tsteps = range(tspan[1], tspan[2], length = datasize)
function trueODEfunc(du, u, p, t)
    true_A = [-0.1 2.0; -2.0 -0.1]
    du .= ((u.^3)'true_A)'
end
prob_trueode = ODEProblem(trueODEfunc, u0, tspan)
# Make the data into a GPU-based array if the user has a GPU
ode_data = gpu(solve(prob_trueode, Tsit5(), saveat = tsteps))


dudt2 = FastChain((x, p) -> x.^3,
                  FastDense(2, 50, tanh),
                  FastDense(50, 2))
u0 = Float32[2.0; 0.0] |> gpu
p = initial_params(dudt2) |> gpu
prob_neuralode = NeuralODE(dudt2, tspan, Tsit5(), saveat = tsteps)

function predict_neuralode(p)
  gpu(prob_neuralode(u0,p))
end
function loss_neuralode(p)
    pred = predict_neuralode(p)
    loss = sum(abs2, ode_data .- pred)
    return loss, pred
end
# Callback function to observe training
list_plots = []
iter = 0
callback = function (p, l, pred; doplot = false)
  global list_plots, iter
  if iter == 0
    list_plots = []
  end
  iter += 1
  display(l)
  # plot current prediction against data
  plt = scatter(tsteps, Array(ode_data[1,:]), label = "data")
  scatter!(plt, tsteps, Array(pred[1,:]), label = "prediction")
  push!(list_plots, plt)
  if doplot
    display(plot(plt))
  end
  return false
end
result_neuralode = DiffEqFlux.sciml_train(loss_neuralode, p,
                                          ADAM(0.05), cb = callback,
                                          maxiters = 300)

---

the error I get is

ArgumentError: cannot take the CPU address of a CuArray{Float32,1}

Stacktrace:
[1] unsafe_convert(::Type{Ptr{Float32}}, ::CuArray{Float32,1}) at /root/.julia/packages/CUDA/dZvbp/src/array.jl:237
[2] unsafe_convert(::Type{Ptr{Float32}}, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}) at ./reshapedarray.jl:292
[3] gemv!(::Char, ::Float32, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Float32, ::CuArray{Float32,1}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/blas.jl:626
[4] gemv!(::CuArray{Float32,1}, ::Char, ::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}, ::Bool, ::Bool) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/matmul.jl:470
[5] mul! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/matmul.jl:66 [inlined]
[6] mul! at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/matmul.jl:208 [inlined]
[7] *(::SubArray{Float32,2,CuArray{Float32,1},Tuple{Base.ReshapedArray{Int64,2,UnitRange{Int64},Tuple{}}},true}, ::CuArray{Float32,1}) at /buildworker/worker/package_linux64/build/usr/share/julia/stdlib/v1.5/LinearAlgebra/src/matmul.jl:47
[8] adjoint at /root/.julia/packages/DiffEqFlux/wifQ9/src/fast_layers.jl:57 [inlined]
[9] pullback at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
[10] applychain at /root/.julia/packages/DiffEqFlux/wifQ9/src/fast_layers.jl:20 [inlined]
[11] pullback(::Zygote.Context, ::typeof(DiffEqFlux.applychain), ::Tuple{FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
... (the last 2 lines are repeated 1 more time)
[14] FastChain at /root/.julia/packages/DiffEqFlux/wifQ9/src/fast_layers.jl:21 [inlined]
[15] pullback(::Zygote.Context, ::FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[16] dudt at /root/.julia/packages/DiffEqFlux/wifQ9/src/neural_de.jl:73 [inlined]
[17] pullback(::Zygote.Context, ::DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}, ::Float32) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[18] pullback at /root/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:38 [inlined]
[19] adjoint at /root/.julia/packages/DiffEqBase/zXRWE/src/zygote.jl:40 [inlined]
[20] pullback at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:47 [inlined]
[21] #30 at /root/.julia/packages/DiffEqSensitivity/GNNnG/src/local_sensitivity/derivative_wrappers.jl:290 [inlined]
[22] pullback(::Zygote.Context, ::DiffEqSensitivity.var"#30#32"{Float32,ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing}}, ::CuArray{Float32,1}, ::CuArray{Float32,1}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[23] pullback at /root/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:38 [inlined]
[24] pullback at /root/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:44 [inlined]
[25] vecjacobian!(::CuArray{Float32,1}, ::CuArray{Float32,1}, ::CuArray{Float32,1}, ::CuArray{Float32,1}, ::Float32, ::DiffEqSensitivity.ODEInterpolatingAdjointSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Base.OneTo{Int64},UnitRange{Int64},LinearAlgebra.UniformScaling{Bool}},InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.Interpolation...
[26] #vecjacobian!#20 at /root/.julia/packages/DiffEqSensitivity/GNNnG/src/local_sensitivity/derivative_wrappers.jl:147 [inlined]
[27] (::DiffEqSensitivity.ODEInterpolatingAdjointSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Base.OneTo{Int64},UnitRange{Int64},LinearAlgebra.UniformScaling{Bool}},InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),D...
[28] ODEFunction at /root/.julia/packages/DiffEqBase/zXRWE/src/diffeqfunction.jl:248 [inlined]
[29] initialize!(::OrdinaryDiffEq.ODEIntegrator{Tsit5,true,CuArray{Float32,1},Nothing,Float32,CuArray{Float32,1},Float32,Float32,Float32,Array{CuArray{Float32,1},1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},true,CuArray{Float32,1},ODEFunction{true,DiffEqSensitivity.ODEInterpolatingAdjointSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Base.OneTo{Int64},UnitRange{Int64},LinearAlgebra.UniformScaling{Bool}},InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{B...
[30] init(::ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},true,CuArray{Float32,1},ODEFunction{true,DiffEqSensitivity.ODEInterpolatingAdjointSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Base.OneTo{Int64},UnitRange{Int64},LinearAlgebra.UniformScaling{Bool}},InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunc...
[31] #solve#391 at /root/.julia/packages/OrdinaryDiffEq/7uPDi/src/solve.jl:4 [inlined]
[32] solve_call(::ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},true,CuArray{Float32,1},ODEFunction{true,DiffEqSensitivity.ODEInterpolatingAdjointSensitivityFunction{DiffEqSensitivity.AdjointDiffCache{Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Base.OneTo{Int64},UnitRange{Int64},LinearAlgebra.UniformScaling{Bool}},InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODE...
[33] #solve_up#458 at /root/.julia/packages/DiffEqBase/zXRWE/src/solve.jl:86 [inlined]
[34] #solve#457 at /root/.julia/packages/DiffEqBase/zXRWE/src/solve.jl:74 [inlined]
[35] adjoint_sensitivities(::ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Fl...
[36] adjoint_sensitivities(::ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Fl...
[37] adjoint_sensitivities(::ODESolution{Float32,2,Array{CuArray{Float32,1},1},Nothing,Nothing,Array{Float32,1},Array{Array{CuArray{Float32,1},1},1},ODEProblem{CuArray{Float32,1},Tuple{Float32,Float32},false,CuArray{Float32,1},ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Float32,Float64,Float64},Tuple{Symbol},NamedTuple{(:saveat,),Tuple{StepRangeLen{Float32,Float64,Float64}}}}}},LinearAlgebra.UniformScaling{Bool},Nothing,typeof(DiffEqFlux.basic_tgrad),Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}},DiffEqBase.StandardODEProblem},Tsit5,OrdinaryDiffEq.InterpolationData{ODEFunction{false,DiffEqFlux.var"#dudt#145"{NeuralODE{FastChain{Tuple{var"#11#12",FastDense{typeof(tanh),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#130"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}},Array{Float32,1},Nothing,Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,StepRangeLen{Flo...
[38] (::DiffEqSensitivity.var"#adjoint_sensitivity_backpass#133"{Tsit5,InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},CuArray{Float32,1},Tuple{},Colon})(::CuArray{Float32,2}) at /root/.julia/packages/DiffEqSensitivity/GNNnG/src/local_sensitivity/concrete_solve.jl:145
[39] #672#back at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:55 [inlined]
[40] #145 at /root/.julia/packages/Zygote/chgvX/src/lib/lib.jl:175 [inlined]
[41] (::Zygote.var"#1681#back#147"{Zygote.var"#145#146"{DiffEqBase.var"#672#back#471"{DiffEqSensitivity.var"#adjoint_sensitivity_backpass#133"{Tsit5,InterpolatingAdjoint{0,true,Val{:central},ZygoteVJP,Bool},CuArray{Float32,1},CuArray{Float32,1},Tuple{},Colon}},Tuple{NTuple{6,Nothing},Tuple{Nothing}}}})(::CuArray{Float32,2}) at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
[42] #solve#457 at /root/.julia/packages/DiffEqBase/zXRWE/src/solve.jl:74 [inlined]
[43] (::typeof(∂(#solve#457)))(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[44] (::Zygote.var"#145#146"{typeof(∂(#solve#457)),Tuple{NTuple{6,Nothing},Tuple{Nothing}}})(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/lib/lib.jl:175
[45] (::Zygote.var"#1681#back#147"{Zygote.var"#145#146"{typeof(∂(#solve#457)),Tuple{NTuple{6,Nothing},Tuple{Nothing}}}})(::CuArray{Float32,2}) at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
[46] (::typeof(∂(solve##kw)))(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[47] (::Zygote.var"#145#146"{typeof(∂(solve##kw)),Tuple{Tuple{Nothing,Nothing,Nothing},Tuple{Nothing}}})(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/lib/lib.jl:175
[48] #1681#back at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49 [inlined]
[49] NeuralODE at /root/.julia/packages/DiffEqFlux/wifQ9/src/neural_de.jl:77 [inlined]
[50] (::typeof(∂(λ)))(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[51] predict_neuralode at ./In[10]:26 [inlined]
[52] (::typeof(∂(predict_neuralode)))(::CuArray{Float32,2}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[53] loss_neuralode at ./In[10]:29 [inlined]
[54] (::typeof(∂(loss_neuralode)))(::Tuple{Float32,Nothing}) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[55] #145 at /root/.julia/packages/Zygote/chgvX/src/lib/lib.jl:175 [inlined]
[56] (::Zygote.var"#1681#back#147"{Zygote.var"#145#146"{typeof(∂(loss_neuralode)),Tuple{Tuple{Nothing},Tuple{}}}})(::Tuple{Float32,Nothing}) at /root/.julia/packages/ZygoteRules/6nssF/src/adjoint.jl:49
[57] #74 at /root/.julia/packages/DiffEqFlux/wifQ9/src/train.jl:120 [inlined]
[58] (::typeof(∂(λ)))(::Float32) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface2.jl:0
[59] (::Zygote.var"#54#55"{Zygote.Params,Zygote.Context,typeof(∂(λ))})(::Float32) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:177
[60] gradient(::Function, ::Zygote.Params) at /root/.julia/packages/Zygote/chgvX/src/compiler/interface.jl:54
[61] macro expansion at /root/.julia/packages/DiffEqFlux/wifQ9/src/train.jl:119 [inlined]
[62] macro expansion at /root/.julia/packages/ProgressLogging/BBN0b/src/ProgressLogging.jl:328 [inlined]
[63] (::DiffEqFlux.var"#73#78"{var"#13#15",Int64,Bool,Bool,typeof(loss_neuralode),CuArray{Float32,1},Zygote.Params})() at /root/.julia/packages/DiffEqFlux/wifQ9/src/train.jl:64
[64] with_logstate(::Function, ::Any) at ./logging.jl:408
[65] with_logger at ./logging.jl:514 [inlined]
[66] maybe_with_logger(::DiffEqFlux.var"#73#78"{var"#13#15",Int64,Bool,Bool,typeof(loss_neuralode),CuArray{Float32,1},Zygote.Params}, ::LoggingExtras.TeeLogger{Tuple{LoggingExtras.EarlyFilteredLogger{ConsoleProgressMonitor.ProgressLogger,DiffEqFlux.var"#68#70"},LoggingExtras.EarlyFilteredLogger{Base.CoreLogging.SimpleLogger,DiffEqFlux.var"#69#71"}}}) at /root/.julia/packages/DiffEqFlux/wifQ9/src/train.jl:39
[67] sciml_train(::Function, ::CuArray{Float32,1}, ::ADAM, ::Base.Iterators.Cycle{Tuple{DiffEqFlux.NullData}}; cb::Function, maxiters::Int64, progress::Bool, save_best::Bool) at /root/.julia/packages/DiffEqFlux/wifQ9/src/train.jl:63
[68] top-level scope at In[10]:52

---

I am using Julia 1.5.2 (same error using 1.5.0)
and library version ...

• BlackBoxOptim v0.5.0
  [a9c8d775] + CPUTime v1.0.0
  [88cd18e8] + ConsoleProgressMonitor v0.1.2
  [459566f4] + DiffEqCallbacks v2.14.1
  [aae7a2af] + DiffEqFlux v1.23.0
  [c894b116] + DiffEqJump v6.10.1
  [77a26b50] + DiffEqNoiseProcess v5.3.1
  [41bf760c] + DiffEqSensitivity v6.31.6
  [ced4e74d] + DistributionsAD v0.6.9
  [a5e1c1ea] + LatinHypercubeSampling v1.6.4
  [1d6d02ad] + LeftChildRightSiblingTrees v0.1.2
  [e6f89c97] + LoggingExtras v0.4.2
  [429524aa] + Optim v1.2.0
  [e409e4f3] + PoissonRandom v0.4.0
  [85a6dd25] + PositiveFactorizations v0.2.3
  [33c8b6b6] + ProgressLogging v0.1.3
  [92933f4c] + ProgressMeter v1.4.0
  [8a4e6c94] + QuasiMonteCarlo v0.2.0
  [e6cf234a] + RandomNumbers v1.4.0
  [ae5879a3] + ResettableStacks v1.0.0
  [37e2e3b7] + ReverseDiff v1.4.3
  [ed01d8cd] + Sobol v1.4.0
  [d4ead438] + SpatialIndexing v0.1.2
  [789caeaf] + StochasticDiffEq v6.26.0
  [5d786b92] + TerminalLoggers v0.1.2
  [9f7883ad] + Tracker v0.2.12

[ChrisRackauckas]

This issue seems to be a bug in CUDA.jl (JuliaGPU/CUDA.jl#448) which I can put a workaround in #420 to release and then we can remove the workaround when the GPU library gets fixed.