Cannot mutate OptState from a different trace level

[LIANHAN3920]

Code

class MLP(nnx.Module):
  def __init__(self, din, dhidden, dout, *, rngs: nnx.Rngs):
    self.linear1 = nnx.Linear(din, dhidden, rngs=rngs)
    self.linear2 = nnx.Linear(dhidden, dout, rngs=rngs)

  def __call__(self, x):
    x = self.linear1(x)
    x = jax.nn.relu(x)
    x = self.linear2(x)
    return x

model = MLP(20,60,10,rngs=nnx.Rngs(42))

optimizer = nnx.Optimizer(model, optax.adam(1e-3))
graph, params = nnx.split(model)

# Define the loss function
def loss_fn(params, batch):
    X_batch, y_batch = batch
    model = nnx.merge(graph,params)
    preds = model(X_batch)
    loss = jnp.mean((preds - y_batch) ** 2)
    return loss

@jit
def train_step(params, batch):
    grads_fn = nnx.value_and_grad(loss_fn)
    loss, grads = grads_fn(params, batch)
    optimizer.update(grads=grads)  <--------------------where erroe occurs
    return params, loss

@jit
def eval_step(params, batch):
    loss = loss_fn(params, batch)
    return loss

total_steps = 10_000
for step, batch in enumerate(dataset(32)):
    print(step)
    params, loss = train_step(params, batch)

    if step % 1000 == 0:
        logs = eval_step(params, batch)
        print(f"step: {step}, loss: {logs['loss']}")

    if step >= total_steps - 1:
        break

model = nnx.merge(graph, params,)
nnx.display(model)

Error

flax.nnx.nnx.errors.TraceContextError: Cannot mutate OptState from a different trace level

not sure if it is because self called in optimizer.update funtion