fix import

stanford-crfm · Apr 24, 2024 · df56c54 · df56c54
1 parent 0dc998a
commit df56c54
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/levanter/optim/util.py b/src/levanter/optim/util.py
@@ -1,5 +1,6 @@
 import equinox as eqx
 import jax
+import jax.numpy as jnp
 
 from levanter.utils.jax_utils import is_inexact_arrayish