Fix #4497: Enable feature importance property for DART booster (#4525)

dmlc · May 31, 2019 · c2a3902 · c2a3902
1 parent ea44417
commit c2a3902
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python-package/xgboost/sklearn.py b/python-package/xgboost/sklearn.py
@@ -536,7 +536,7 @@ def feature_importances_(self):
         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()