I don’t know enough about ODEs to suggest specific optimizations right now, but the main problem (as you’ve noticed) is that performing the operations increases the size of the computation graph. Each addition/multiplication is an operation in the computation graph, and it sounds like the problem is that the computation graph is getting too large. Do you notice your memory usage steadily growing with each iteration of your for loop?
We’re working on new ways to trade compute for memory (see this pull request 15) that will probably help in this case.