aux

kostastsa · Nov 1, 2023 · edbfb9c · edbfb9c
1 parent 10a7860
commit edbfb9c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dynamax/ssm.py b/dynamax/ssm.py
@@ -4,7 +4,7 @@
 from functools import partial
 import jax.numpy as jnp
 import jax.random as jr
-from jax import jit, lax, vmap, debug
+from jax import jit, lax, vmap
 from jax.tree_util import tree_map
 from jaxtyping import Float, Array, PyTree
 import optax