diff --git a/autogalaxy/quantity/model/visualizer.py b/autogalaxy/quantity/model/visualizer.py
index 32483aab..e3bfbbbd 100644
--- a/autogalaxy/quantity/model/visualizer.py
+++ b/autogalaxy/quantity/model/visualizer.py
@@ -1,6 +1,7 @@
 import os
 
 import autofit as af
+from autofit.non_linear.test_mode import is_test_mode
 
 from autogalaxy.quantity.model.plotter import PlotterQuantity
 
@@ -65,7 +66,7 @@ def visualize(
             via a non-linear search).
         """
 
-        if os.environ.get("PYAUTOFIT_TEST_MODE") == "1":
+        if is_test_mode():
             return
 
         fit = analysis.fit_quantity_for_instance(instance=instance)