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

This commit is contained in:
Philip Hyunsu Cho 2019-05-31 15:11:57 -07:00 committed by GitHub
parent ea44417754
commit c2a3902ba3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -536,7 +536,7 @@ class XGBModel(XGBModelBase):
feature_importances_ : array of shape ``[n_features]`` 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 {}' raise AttributeError('Feature importance is not defined for Booster type {}'
.format(self.booster)) .format(self.booster))
b = self.get_booster() b = self.get_booster()