diff --git a/src/concrete/ml/sklearn/base.py b/src/concrete/ml/sklearn/base.py
index 6950a0af0..c6355de16 100644
--- a/src/concrete/ml/sklearn/base.py
+++ b/src/concrete/ml/sklearn/base.py
@@ -808,11 +808,9 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
             # If the prediction array is 1D, transform the output into a 2D array [1-p, p],
             # with p the initial output probabilities
             # This is similar to what is done in scikit-learn
-            if y_preds.ndim == 1:
-                y_preds = numpy.vstack([1 - y_preds, y_preds]).T
-
-            elif y_preds.shape[1] == 1:
-                y_preds = numpy.concatenate((1 - y_preds, y_preds), axis=1)
+            if y_preds.ndim == 1 or y_preds.shape[1] == 1:
+                y_preds = y_preds.flatten()
+                return numpy.vstack([1 - y_preds, y_preds]).T
 
         # Else, apply the softmax operator
         else:
diff --git a/src/concrete/ml/sklearn/linear_model.py b/src/concrete/ml/sklearn/linear_model.py
index 6e809f441..afd02991d 100644
--- a/src/concrete/ml/sklearn/linear_model.py
+++ b/src/concrete/ml/sklearn/linear_model.py
@@ -868,9 +868,6 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
             assert isinstance(self.classes_, numpy.ndarray)
             binary = len(self.classes_) == 2
 
-            if binary:
-                y_preds = y_preds[:, 0]
-
             prob2 = numpy.empty(tuple())
             if binary:
                 prob2 = numpy.ones((y_preds.shape[0], 2))