set matplotlib backend (#785)

mljar · Jan 10, 2025 · 1508365 · 1508365
1 parent da2aa81
commit 1508365
Showing 1 changed file with 8 additions and 2 deletions.
diff --git 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:
         """