ich

optax.MultiSteps out of memory

15

I always get an out of memory error using optax.MultiSteps, even when every_k_schedule=1. Using optax.apply_every(k=1) in a chain works fine. ``` optimizer = optax.chain( optax.clip_by_global_norm(0.5), optax.adam(lr), #optax.apply_every(k=1) ) optimizer =...