From 4fac9874e0abcf864a9ca00f66ccb65bcdd86696 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Tue, 22 Jan 2019 16:03:54 +0800 Subject: [PATCH] Check booster for dart in feature importance. (#4073) * Check booster for dart in feature importance. --- python-package/xgboost/core.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/python-package/xgboost/core.py b/python-package/xgboost/core.py index cde339bb8..3bb7299ea 100644 --- a/python-package/xgboost/core.py +++ b/python-package/xgboost/core.py @@ -1448,8 +1448,7 @@ class Booster(object): importance_type: str, default 'weight' One of the importance types defined above. """ - - 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 ValueError('Feature importance is not defined for Booster type {}' .format(self.booster))