diff options
| -rw-r--r-- | lib/parameters.py | 6 |
1 files changed, 5 insertions, 1 deletions
diff --git a/lib/parameters.py b/lib/parameters.py index 0c4401c..58af96f 100644 --- a/lib/parameters.py +++ b/lib/parameters.py @@ -611,7 +611,11 @@ class ModelAttribute: return f"ModelAttribute<{self.name}, {self.attr}, mean={mean}>" def to_json(self, **kwargs): - if type(self.model_function) in (df.CARTFunction, df.XGBoostFunction): + if type(self.model_function) in ( + df.CARTFunction, + df.LMTFunction, + df.XGBoostFunction, + ): import sklearn.tree feature_names = list( |