# Build an abstract model, no real memory allocated
abstract_model = nnx.eval_shape(lambda: CNN(rngs=nnx.Rngs(0)))

# Split into structure and abstract state
graphdef, abstract_state = nnx.split(abstract_model)

# Ask the manager for the latest available step
step_to_restore = mngr.latest_step()

if step_to_restore is not None:
    restored_state = mngr.restore(
        step_to_restore,
        args=ocp.args.StandardRestore(abstract_state),
    )
    
    # 4. Reconstruct a live, usable model
    restored_model = nnx.merge(graphdef, restored_state)

mngr.close()

        __
        __