Merge 2ab5517 into 43a189c

BayesWitnesses · Jan 19, 2020 · af56878 · af56878
2 parents 43a189c + 2ab5517
commit af56878
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/m2cgen/assemblers/boosting.py b/m2cgen/assemblers/boosting.py
@@ -206,8 +206,12 @@ def __init__(self, model,
                  leaves_cutoff_threshold=LEAVES_CUTOFF_THRESHOLD):
         model_dump = model.booster_.dump_model()
         trees = [m["tree_structure"] for m in model_dump["tree_info"]]
+
         self.n_iter = len(trees) // model_dump["num_tree_per_iteration"]
-        self.average_output = model_dump["average_output"]
+        if "average_output" in model_dump:
+            self.average_output = model_dump["average_output"]
+        else:
+            self.average_output = False
 
         super().__init__(model, trees,
                          leaves_cutoff_threshold=leaves_cutoff_threshold)