Update rigid_body_test.py

abhijeetgangan · May 22, 2024 · 36691c3 · 36691c3
1 parent e20bffa
commit 36691c3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/rigid_body_test.py b/tests/rigid_body_test.py
@@ -29,7 +29,6 @@
 from jax import test_util as jtu
 
 from jax import config as jax_config
-jax_config.update('jax_disable_jit', True)
 
 import jax.numpy as jnp
 
@@ -45,6 +44,7 @@
 
 from functools import partial
 
+jax_config.update('jax_disable_jit', True)
 jax_config.parse_flags_with_absl()