diff --git a/autolens/analysis/analysis/dataset.py b/autolens/analysis/analysis/dataset.py
index a74b9d300..41bb80847 100644
--- a/autolens/analysis/analysis/dataset.py
+++ b/autolens/analysis/analysis/dataset.py
@@ -84,6 +84,11 @@ def __init__(
             anyway.
         """
 
+        import os
+
+        if os.environ.get("PYAUTO_DISABLE_JAX") == "1":
+            use_jax = False
+
         super().__init__(
             dataset=dataset,
             adapt_images=adapt_images,
@@ -108,9 +113,6 @@ def __init__(
         if is_test_mode():
             self.raise_inversion_positions_likelihood_exception = False
 
-        # Can be deleted after relevent AutoFIT PR merged
-        self._use_jax = use_jax
-
     def modify_before_fit(self, paths: af.DirectoryPaths, model: af.Collection):
         """
         This function is called immediately before the non-linear search begins and performs final tasks and checks