Improve state error

patrick-kidger · Sep 14, 2024 · 5f588e5 · 5f588e5
1 parent ec4747e
commit 5f588e5
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/equinox/nn/_stateful.py b/equinox/nn/_stateful.py
@@ -119,9 +119,8 @@ def __init__(self, model: PyTree):
             if _is_index(leaf):
                 if leaf.init is _sentinel:
                     raise ValueError(
-                        "Cannot call `eqx.nn.State(eqx.nn.delete_init_state(model))`. "
-                        "You should call `eqx.nn.State(model)`, using the original "
-                        "model."
+                        "Do not call `eqx.nn.State(model)` directly. You should call "
+                        "`eqx.nn.make_with_state(ModelClass)(...args...)` instead."
                     )
                 state[leaf.marker] = jtu.tree_map(jnp.asarray, leaf.init)
         self._state = state