diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index 423de707..415b63a3 100644 --- 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 diff --git a/tests/lightgbm/test_objective_functions.py b/tests/lightgbm/test_objective_functions.py index 9fc2db8a..b9bd2f91 100644 --- a/tests/lightgbm/test_objective_functions.py +++ b/tests/lightgbm/test_objective_functions.py @@ -31,6 +31,7 @@ class ObjectiveTest(unittest.TestCase): "regression", "poisson", "gamma", + "quantile" ) @staticmethod