From d2a7296acd18a05f3754453f3009380120634e0c Mon Sep 17 00:00:00 2001
From: Abhijeet Gangan <35770546+abhijeetgangan@users.noreply.github.com>
Date: Tue, 21 May 2024 22:35:29 -0700
Subject: [PATCH] Revert "Try #2"

---
 tests/rigid_body_test.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/rigid_body_test.py b/tests/rigid_body_test.py
index 7316f5d..fd59675 100644
--- a/tests/rigid_body_test.py
+++ b/tests/rigid_body_test.py
@@ -28,10 +28,7 @@
 from jax import lax
 from jax import test_util as jtu
 
-
-from jax.config import config as jax_config
-jax_config.update('jax_disable_jit', True)
-
+from jax import config as jax_config
 import jax.numpy as jnp
 
 from jax_md import quantity