Augmented neural ODE on GPU

Specific Domains Machine Learning
, , ,
1

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?

Thanks

1 Like
2

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

3

Sure thing! Here it is.

1 Like