Nested gradient computations with differential equation solver

In the code above only 1 parameter but my end goal would be to train some models with a few hundred thousand parameters.
I’ve played around with Flux.jl, and this is why went with Zygote.jl by default.
For your information, I am just starting out in Julia and have been using pytorch so far therefore I haven’t had to think too much about the AD part of my work so all your comments are very much appreciated :wink: