Augmented neural ODE on GPU

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?


1 Like

@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.

Sure thing! Here it is.

1 Like