From 150836538f1188f1452247d2e6519529a2f4e903 Mon Sep 17 00:00:00 2001
From: Piotr <pplonski86@gmail.com>
Date: Fri, 10 Jan 2025 14:45:11 +0100
Subject: [PATCH] set matplotlib backend (#785)

---
 supervised/automl.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

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:
         """