Fix future warning by avoiding importing of private modules (#200)

BayesWitnesses · May 5, 2020 · 0b2a93d · 0b2a93d
1 parent 91dbb6e
commit 0b2a93d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tests/utils.py b/tests/utils.py
@@ -13,9 +13,9 @@
 from lightning.impl.base import BaseClassifier as LightBaseClassifier
 from sklearn import datasets
 from sklearn.base import BaseEstimator, RegressorMixin, clone
-from sklearn.ensemble import forest
+from sklearn.ensemble._forest import ForestClassifier
 from sklearn.utils import shuffle
-from sklearn.linear_model.base import LinearClassifierMixin
+from sklearn.linear_model._base import LinearClassifierMixin
 from sklearn.tree import DecisionTreeClassifier
 from sklearn.svm import SVC, NuSVC
 from xgboost import XGBClassifier
@@ -129,7 +129,7 @@ def __call__(self, estimator):
             y_pred = estimator.predict_proba(self.X_test.astype(np.float32))
         elif isinstance(
                 estimator,
-                (forest.ForestClassifier, XGBClassifier, LGBMClassifier)):
+                (ForestClassifier, XGBClassifier, LGBMClassifier)):
             y_pred = estimator.predict_proba(self.X_test)
         else:
             y_pred = estimator.predict(self.X_test)