add support for quantile objective for LGBM models

Signed-off-by: Mohammed Ajil <mohammed@ajil.ch>
onnx · Sep 23, 2021 · 14c2691 · 14c2691
1 parent 2351850
commit 14c2691
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py
@@ -444,7 +444,7 @@ def convert_lightgbm(scope, operator, container):
     elif gbm_text['objective'].startswith('multiclass'):
         n_classes = gbm_text['num_class']
         attrs['post_transform'] = 'SOFTMAX'
-    elif gbm_text['objective'].startswith('regression'):
+    elif gbm_text['objective'].startswith(('regression', 'quantile')):
         n_classes = 1  # Regressor has only one output variable
         attrs['post_transform'] = 'NONE'
         attrs['n_targets'] = n_classes