Fix sklearn regressor test?

mgq1507 · Aug 22, 2019 · fb7f49e · fb7f49e
1 parent c074416
commit fb7f49e
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/keras/wrappers/scikit_learn.py b/keras/wrappers/scikit_learn.py
@@ -320,8 +320,10 @@ def predict(self, x, **kwargs):
                 Predictions.
         """
         kwargs = self.filter_sk_params(Sequential.predict, kwargs)
-        preds = self.model.predict(x, **kwargs)
-        return np.squeeze(preds, axis=len(preds.shape) - 1)
+        preds = np.array(self.model.predict(x, **kwargs))
+        if preds.shape[-1] == 1:
+            return np.squeeze(preds, axis=-1)
+        return preds
 
     def score(self, x, y, **kwargs):
         """Returns the mean loss on the given test data and labels.