code: def create_and_save(seed: int, path: str):
      model = create_model(seed)
      state = model.get_state()
      # Save the parameters
      checkpointer = orbax.PyTreeCheckpointer()
      checkpointer.save(f"{path}/state", state)
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/08_save_load_checkpoints.py
imports: from tempfile import TemporaryDirectory
  import jax
  import jax.numpy
  import orbax.checkpoint
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling a global function with example usage: create_and_save(seed, path) without return types