diff --git a/supervised/automl.py b/supervised/automl.py
index 7b8a535b..03db029a 100644
--- a/supervised/automl.py
+++ b/supervised/automl.py
@@ -2,7 +2,6 @@
 
 import matplotlib
 
-matplotlib.use("Agg")
 import warnings
 
 warnings.filterwarnings("ignore", message=".*The 'nopython' keyword.*")
@@ -429,7 +428,14 @@ def fit(
         Returns:
             AutoML object: Returns `self`
         """
-        return self._fit(X, y, sample_weight, cv, sensitive_features)
+        try:
+            original_backend = matplotlib.get_backend()
+            matplotlib.use("Agg")
+            return self._fit(X, y, sample_weight, cv, sensitive_features)
+        except Exception as e:
+            raise e
+        finally:
+            matplotlib.use(original_backend)
 
     def predict(self, X: Union[List, numpy.ndarray, pandas.DataFrame]) -> numpy.ndarray:
         """