add 'poisson' support also to WrappedBooster

onnx · May 12, 2021 · 3cccfd3 · 3cccfd3
1 parent 5929839
commit 3cccfd3
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py
@@ -27,7 +27,8 @@ def __init__(self, booster):
         if (_model_dict['objective'].startswith('binary') or
                 _model_dict['objective'].startswith('multiclass')):
             self.operator_name = 'LgbmClassifier'
-        elif _model_dict['objective'].startswith('regression'):
+        elif (_model_dict['objective'].startswith('regression') or
+                _model_dict['objective'].startswith('poisson')):
             self.operator_name = 'LgbmRegressor'
         else:
             # Other objectives are not supported.
@@ -170,4 +171,4 @@ def parse_lightgbm(model, initial_types=None, target_opset=None,
     for variable in outputs:
         raw_model_container.add_output(variable)
 
-    return topology
+    return topology