devjalx

Equinox has spikes in training step durations due to garbage collector

15

Hi all, this is related to https://github.com/jax-ml/jax/issues/30185, however, the issue occurs when running already jitted functions in combination with equinox and not during compilation. When training neural networks in equinox,...