Trying to run the augumented neural ODE example from the docs and it works great on CPU. However, moving all the parts with the
gpu function like when using a regular Flux model does not seem to do the trick.
I am getting errors like
MethodError: no method matching AugmentedNDELayer(::NeuralODE...) when calling
gpu on the constructed model.
Has anyone here managed to do this and could give some advice?