Merge f085588 into d73048e

BayesWitnesses · May 5, 2020 · 330865f · 330865f
2 parents d73048e + f085588
commit 330865f
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
@@ -120,7 +120,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)