fix state update in NeuralODE example

[gabrevaya]

I think this is a typo, isn't it? The state should be updated across the training loops, right?

[codecov]

Codecov Report

Merging #107 (70822de) into main (6babea5) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main     #107   +/-   ##
=======================================
  Coverage   76.79%   76.79%           
=======================================
  Files          15       15           
  Lines        1056     1056           
=======================================
  Hits          811      811           
  Misses        245      245           

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6babea5...70822de. Read the comment docs.