diff --git a/onnxmltools/convert/lightgbm/_parse.py b/onnxmltools/convert/lightgbm/_parse.py index e527faf7..d711abe0 100644 --- a/onnxmltools/convert/lightgbm/_parse.py +++ b/onnxmltools/convert/lightgbm/_parse.py @@ -27,9 +27,7 @@ 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') or - _model_dict['objective'].startswith('poisson') or - _model_dict['objective'].startswith('gamma')): + elif _model_dict['objective'].startswith(('regression', 'poisson', 'gamma')): self.operator_name = 'LgbmRegressor' else: # Other objectives are not supported. diff --git a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py index c2ad360c..242c2c1d 100644 --- a/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py +++ b/onnxmltools/convert/lightgbm/operator_converters/LightGbm.py @@ -234,7 +234,7 @@ def convert_lightgbm(scope, operator, container): n_classes = 1 # Regressor has only one output variable attrs['post_transform'] = 'NONE' attrs['n_targets'] = n_classes - elif gbm_text['objective'].startswith('poisson') or gbm_text['objective'].startswith('gamma'): + elif gbm_text['objective'].startswith(('poisson', 'gamma')): n_classes = 1 # Regressor has only one output variable attrs['n_targets'] = n_classes # 'Exp' is not a supported post_transform value in the ONNX spec yet,