improved

README.md

@@ -51,6 +51,8 @@ weights_path = hf_hub_download(repo_id="Jorgvt/ppnet-bio-fitted",
                                filename="weights.safetensors")
 variables = load_file(weights_path)
 variables = flax.traverse_util.unflatten_dict(variables, sep=".")
+state = variables["state"]
+params = variables["params"]
 ```
 
 #### 3.2. Using `mgspack`
@@ -59,10 +61,9 @@ weights_path = hf_hub_download(repo_id="Jorgvt/ppnet-bio-fitted",
                                filename="weights.msgpack")
 with open(weights_path, "rb") as f:
     variables = orbax.checkpoint.msgpack_utils.msgpack_restore(f.read())
+variables = jax.tree_util.tree_map(lambda x: jnp.array(x), variables)
 state = variables["state"]
 params = variables["params"]
-state = jax.tree_util.tree_map(lambda x: jnp.array(x), state)
-params = jax.tree_util.tree_map(lambda x: jnp.array(x), params)
 ```
 
 ### 4. Use the model