diff --git a/python-package/xgboost/sklearn.py b/python-package/xgboost/sklearn.py index ae4751ea7..b80a79fcf 100644 --- a/python-package/xgboost/sklearn.py +++ b/python-package/xgboost/sklearn.py @@ -536,7 +536,7 @@ class XGBModel(XGBModelBase): feature_importances_ : array of shape ``[n_features]`` """ - if getattr(self, 'booster', None) is not None and self.booster != 'gbtree': + if getattr(self, 'booster', None) is not None and self.booster not in {'gbtree', 'dart'}: raise AttributeError('Feature importance is not defined for Booster type {}' .format(self.booster)) b = self.get_booster()