Fix support of older LightGBM versions (#150)

BayesWitnesses · Jan 19, 2020 · 6802632 · 6802632
1 parent 43a189c
commit 6802632
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/m2cgen/assemblers/boosting.py b/m2cgen/assemblers/boosting.py
@@ -206,8 +206,9 @@ 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"]
+        self.average_output = model_dump.get("average_output", False)
 
         super().__init__(model, trees,
                          leaves_cutoff_threshold=leaves_cutoff_threshold)