Augmented neural ODE on GPU errors out #329
Description
Trying to run the ANODE example in the docs on the GPU by moving
data, labels, model, and parameters to GPU using the gpu function gives ArgumentError: cannot take the CPU address of a CuArrays.CuArray{Float32,2,Nothing} when passing data to the model.
Rewriting the code to use "standard" Flux constructs instead of FastChain, FastDense etc. Gives the error AugmentedNDELayer(::NeuralODE{Chain{Tuple{Dense{typeof(relu),Array{Float32,2},Array{Float32,1}},Dense{typeof(relu),Array{Float32,2},Array{Float32,1}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}}},Array{Float32,1},Flux.var"#12#14"{Chain{Tuple{Dense{typeof(relu),Array{Float32,2},Array{Float32,1}},Dense{typeof(relu),Array{Float32,2},Array{Float32,1}},Dense{typeof(identity),Array{Float32,2},Array{Float32,1}}}}},Tuple{Float32,Float32},Tuple{Tsit5},Base.Iterators.Pairs{Symbol,Real,NTuple{4,Symbol},NamedTuple{(:save_everystep, :reltol, :abstol, :save_start),Tuple{Bool,Float64,Float64,Bool}}}}) when calling gpu on the model.
Is there a way to run an ANODE on the GPU? I could very well be missing something.