diff --git a/autolens/analysis/analysis/dataset.py b/autolens/analysis/analysis/dataset.py
index 82d107caf..a74b9d300 100644
--- a/autolens/analysis/analysis/dataset.py
+++ b/autolens/analysis/analysis/dataset.py
@@ -31,7 +31,7 @@
 from autolens.analysis.positions import PositionsLH
 
 from autolens import exc
-from autofit.non_linear.test_mode import is_test_mode
+from autoconf.test_mode import is_test_mode
 
 logger = logging.getLogger(__name__)
 
diff --git a/autolens/analysis/result.py b/autolens/analysis/result.py
index f147d2cb3..2a7ac3fe5 100644
--- a/autolens/analysis/result.py
+++ b/autolens/analysis/result.py
@@ -31,7 +31,7 @@
 )
 from autolens.lens.tracer import Tracer
 from autolens.point.solver import PointSolver
-from autofit.non_linear.test_mode import is_test_mode
+from autoconf.test_mode import is_test_mode
 
 logger = logging.getLogger(__name__)
 
diff --git a/autolens/quantity/model/visualizer.py b/autolens/quantity/model/visualizer.py
index f29fd474d..d2b483946 100644
--- a/autolens/quantity/model/visualizer.py
+++ b/autolens/quantity/model/visualizer.py
@@ -1,7 +1,7 @@
 import os
 
 import autofit as af
-from autofit.non_linear.test_mode import is_test_mode
+from autoconf.test_mode import is_test_mode
 
 from autogalaxy.quantity.model.plotter import PlotterQuantity