[python] fixed result shape in case of predict_proba with raw_score arg

microsoft · Nov 25, 2018 · cc99f0d · cc99f0d
1 parent 0c4bb89
commit cc99f0d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/python-package/lightgbm/sklearn.py b/python-package/lightgbm/sklearn.py
@@ -796,7 +796,7 @@ def predict_proba(self, X, raw_score=False, num_iteration=None,
         """
         result = super(LGBMClassifier, self).predict(X, raw_score, num_iteration,
                                                      pred_leaf, pred_contrib, **kwargs)
-        if self._n_classes > 2 or pred_leaf or pred_contrib:
+        if self._n_classes > 2 or raw_score or pred_leaf or pred_contrib:
             return result
         else:
             return np.vstack((1. - result, result)).transpose()