Import matplotlib in log_feature_importance_plot (mlflow#2423)

avflor · Aug 22, 2020 · ef7fea0 · ef7fea0
1 parent f33cc47
commit ef7fea0
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/mlflow/lightgbm.py b/mlflow/lightgbm.py
@@ -212,7 +212,6 @@ def autolog():
     """
     import lightgbm
     import numpy as np
-    import matplotlib.pyplot as plt
 
     @gorilla.patch(lightgbm)
     def train(*args, **kwargs):
@@ -234,6 +233,8 @@ def log_feature_importance_plot(features, importance, importance_type):
             """
             Log feature importance plot.
             """
+            import matplotlib.pyplot as plt
+
             indices = np.argsort(importance)
             features = np.array(features)[indices]
             importance = importance[indices]

diff --git a/mlflow/xgboost.py b/mlflow/xgboost.py
@@ -216,7 +216,6 @@ def autolog(importance_types=['weight']):  # pylint: disable=W0102
     """
     import xgboost
     import numpy as np
-    import matplotlib.pyplot as plt
 
     @gorilla.patch(xgboost)
     def train(*args, **kwargs):
@@ -239,6 +238,8 @@ def log_feature_importance_plot(features, importance, importance_type):
             """
             Log feature importance plot.
             """
+            import matplotlib.pyplot as plt
+
             features = np.array(features)
             importance = np.array(importance)
             indices = np.argsort(importance)