Axon memory usage high

,

Ah right. I thought that was the whole Nx code.

You could add a Stream.map function that calls :erlang.gc before every epoch at least, and that will eliminate the possibility of gc being too slow.

Also, which EXLA version are you using?