I’ve continued investigating this and I have a little bit more of insight. The first forward pass of the UDE works fine, the initial values of the NN are stable and produce a meaningful loss. However, right after, the ODE solver is called again (I can see the progress bar), but the new initial conditions (i.e. H) are exactly the same as at the end of the previous forward run. And then is when the gradients go to infinity and everything crashes.
Since sciml_train is such a high level black box, I’m having a hard time understanding what is going on exactly. Is that the pullback calling the solver backwards? Is that why the initial conditions match the final conditions of the previous forward run? Not sure if this is of much help, just some extra context that might avoid you trying to debug it. Thanks again!