Question about the optimization w.r.t. a process

Just as mentioned in the “Heightfield fluid simulation” example during lec01 of GAMES201, an optimization is performed w.r.t. an objective defined over a process (physical process consisting of 256 timesteps in this case).

My question is: what is the requirement for the computation at each time step? E.g. should it be able to be rewritten in Taichi or it accepts computation out of the scope of Taichi (e.g. a forward of an NN in PyTorch)?

Besides, which factors will influence the memory footprint (memory consumption)?