ODE solver getting stuck for simple term

Hey, I am trying to implement a continuous normalizing flow, where the vector field is parameterized as the gradient of a potential function. Depending on the Adjoint method, either I am getting this error (for the default adjoint method):

ValueError: Reverse-mode differentiation does not work for lax.while_loop or lax.fori_loop. Try using lax.scan instead.

or, for adjoint=diffrax.BacksolveAdjoint(), the ODE solver runs for some epochs and then suddenly gets stuck without any error messages (when digging into the debugger, it seems like as soon as the solver approaches t=0, it restarts over and over again), and even after letting it run for an hour it didn't make any progress (whereas the epochs before took some seconds maximally.)
I tried different solvers (Dopri5, Kvaerno5(nonlinear_solver=NewtonNonlinearSolver())), plotted the vector field to make sure it is well behaving (it seems to be) and enabled float64, but nothing helped.
If I parameterize the vector field directly as a vector-valued output of the NN it works for adjoint=diffrax.BacksolveAdjoint() (but the same error as above for the default adjoint method).
I ported the code into a google colab here
Do you have any idea what the reasons might be for this behaviour?

Start with the linked Google Colab and reproduce the continuous normalizing flow with the default adjoint and BacksolveAdjoint methods. Compare the Dopri5 and Kvaerno5 runs as they approach t=0, then determine whether the behavior is caused by the gradient-parameterized vector field or the adjoint method. Done means documenting an actionable cause, workaround, or confirmed solver bug.