fix test

sail-sg · Jan 12, 2024 · 7587f91 · 7587f91
1 parent 37be812
commit 7587f91
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/autofd/general_array_test.py b/autofd/general_array_test.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import optax
 import jax
 import jax.numpy as jnp
 from jaxtyping import Float32, Array