optax.MultiSteps out of memory

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 = optax.MultiSteps(optimizer, every_k_schedule=1)
```
Later I'm using
`opt_state = optimizer.init(params)`
and
```
updates, opt_state = optimizer.update(grads, opt_state, params)
params = optax.apply_updates(params, updates)
```
I have no idea what I could be doing wrong. I'm not changing anything else, like batch size.