existing_model = CNN(rngs=nnx.Rngs(0))
nnx.update(existing_model, restored_state)

        __
        __