How to use optax.ema?

[noeltong]

Hello, everyone! I'm just training my model with adamw optimizer and exponential moving average, and I used optax.chain to combine the gradient clip, the optimizer and optax.ema,

optimizer = optax.chain(
            optax.adaptive_grad_clip(1.),
            optimizer,
            optax.ema(0.9999)
        )

However, when I started training my model, the loss seems to be hard to converge

Did I use optax.ema properly? The version of optax is 0.1.3.