minor change to return

This commit is contained in:
Mark Regan 2022-10-27 12:41:12 +01:00
parent 1c98640129
commit 6ef82dd8b6

View File

@ -97,8 +97,10 @@ class FreqaiMultiOutputClassifier(MultiOutputClassifier):
``n_classes``) for that particular output. ``n_classes``) for that particular output.
""" """
check_is_fitted(self) check_is_fitted(self)
results = np.hstack([estimator.predict_proba(X) for estimator in self.estimators_]) results = np.squeeze(np.hstack(
return np.squeeze(results) [estimator.predict_proba(X) for estimator in self.estimators_]
))
return results
def predict(self, X): def predict(self, X):
"""Predict multi-output variable using model for each target variable. """Predict multi-output variable using model for each target variable.
@ -122,6 +124,6 @@ class FreqaiMultiOutputClassifier(MultiOutputClassifier):
delayed(e.predict)(X) for e in self.estimators_ delayed(e.predict)(X) for e in self.estimators_
) )
results = np.asarray(y).T results = np.squeeze(np.asarray(y).T)
return np.squeeze(results) return results