Updated jax.config import

PiperOrigin-RevId: 574896308
google · Oct 19, 2023 · 557f43e · 557f43e
1 parent ee3bf7a
commit 557f43e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/hessian/test_run_stats.py b/hessian/test_run_stats.py
@@ -34,7 +34,7 @@
 from init2winit.init_lib import initializers
 from init2winit.model_lib import models
 from init2winit.trainer_lib import trainer
-from jax.config import config as jax_config
+from jax import config as jax_config
 from jax.flatten_util import ravel_pytree
 import jax.numpy as jnp
 import jax.random