# Augmented neural ODE on GPU

**URL:** <https://discourse.julialang.org/t/augmented-neural-ode-on-gpu/42096>\
**Category:** Machine Learning\
**Tags:** question, diffeq, flux, ode\
**Created:** [June 26, 2020, 11:57am UTC](https://discourse.julialang.org/t/augmented-neural-ode-on-gpu/42096 "2020-06-26T11:57:53Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![SebastianCallh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sebastiancallh/32/16102_2.png) [@SebastianCallh](https://discourse.julialang.org/u/SebastianCallh)\
**Post date:** [June 26, 2020, 11:57am UTC](https://discourse.julialang.org/t/augmented-neural-ode-on-gpu/42096/1 "2020-06-26T11:57:53Z")

</div>

Trying to run the augumented neural ODE example from [the docs](https://diffeqflux.sciml.ai/dev/examples/augmented_neural_ode/) 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?

Thanks

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [June 26, 2020, 1:44pm UTC](https://discourse.julialang.org/t/augmented-neural-ode-on-gpu/42096/2 "2020-06-26T13:44:17Z")

</div>

@avikpal might want to look into this. He just added the augmentation functions last week, so this just hasn’t been handled yet. But it’s a good next step. Can you open an issue on the repo? We can take it from there.

---

<div class="post-metadata">

**Author:** ![SebastianCallh](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sebastiancallh/32/16102_2.png) [@SebastianCallh](https://discourse.julialang.org/u/SebastianCallh)\
**Post date:** [June 26, 2020, 2:07pm UTC](https://discourse.julialang.org/t/augmented-neural-ode-on-gpu/42096/3 "2020-06-26T14:07:14Z")

</div>

Sure thing! [Here](https://github.com/SciML/DiffEqFlux.jl/issues/329) it is.