Update loader

load_model.py

@@ -12,5 +12,5 @@ def load_model(repo_id="Extropic-AI/Z1T-0"):
     weights_path = hf_hub_download(repo_id, "model.eqx")
     with open(config_path) as f:
         config = Config(**json.load(f))
-    model = create_model(config, jax.random.PRNGKey(0))
+    model = create_model(config, jax.random.key(0))
     return eqx.tree_deserialise_leaves(weights_path, model)