Custom NeuralODE layer trains innefficiently

Nope, I created the issue