# Buggy code
def predict_proba(self, X):
    """Predict class probabilities for X.
    The predicted class probabilities of an input sample are computed as
    the mean predicted class probabilities of the trees in the forest. The
    class probability of a single tree is the fraction of samples of the same
    class in a leaf.
    Parameters
    ----------
    X : array-like or sparse matrix of shape = [n_samples, n_features]
        The input samples. Internally, its dtype will be converted to
        ``dtype=np.float32``. If a sparse matrix is provided, it will be
        converted into a sparse ``csr_matrix``.
    Returns
    -------
    p : array of shape = [n_samples, n_classes], or a list of n_outputs
        such arrays if n_outputs > 1.
        The class probabilities of the input samples. The order of the
        classes corresponds to that in the attribute `classes_`.
    """
    check_is_fitted(self, "estimators_")

    # Check data
    if self.check_input:
        X = self._validate_X_predict(X)

    # Assign chunk of trees to jobs
    n_jobs, _, _ = _partition_estimators(self.n_estimators, self.n_jobs)

    all_proba = Parallel(n_jobs=n_jobs, verbose=self.verbose)(
        delayed(e.predict_proba)(X) for e in self.estimators_
    )

    all_proba = np.sum(all_proba, axis=0) / len(self.estimators_)

    if len(all_proba) == 1:
        return all_proba[0]
    else:
        return all_proba


# Traceback (AxisError)
---------------------------------------------------------------------------
AxisError                                 Traceback (most recent call last)
<ipython-input-56-a8b067ca3dcd> in <module>
----> 1 clf.predict(X_test.iloc[0:1, :])

~/.conda/envs/py3/lib/python3.7/site-packages/sklearn/utils/metaestimators.py in <lambda>(*args, **kwargs)
114
115         # lambda, but not partial, allows help() to work with update_wrapper
--> 116         out = lambda *args, **kwargs: self.fn(obj, *args, **kwargs)
117         # update the docstring of the returned function
118         update_wrapper(out, self.fn)

~/.conda/envs/py3/lib/python3.7/site-packages/sklearn/pipeline.py in predict(self, X, **predict_params)
420         for _, name, transform in self._iter(with_final=False):
421             Xt = transform.transform(Xt)
--> 422         return self.steps[-1][-1].predict(Xt, **predict_params)
423
424     @if_delegate_has_method(delegate='_final_estimator')

~/.conda/envs/py3/lib/python3.7/site-packages/sklearn/ensemble/forest.py in predict(self, X)
546
547         if self.n_outputs_ == 1:
--> 548             return self.classes_.take(np.argmax(proba, axis=1), axis=0)
549
550         else:

~/.conda/envs/py3/lib/python3.7/site-packages/numpy/core/fromnumeric.py in argmax(a, axis, out)
1101
1102     """
-> 1103     return _wrapfunc(a, 'argmax', axis=axis, out=out)
1104
1105

~/.conda/envs/py3/lib/python3.7/site-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
54 def _wrapfunc(obj, method, *args, **kwds):
55     try:
---> 56         return getattr(obj, method)(*args, **kwds)
57
58     # An AttributeError occurs if the object does not have

AxisError: axis 1 is out of bounds for array of dimension 1

# Fixed code
def predict_proba(self, X):
    """Predict class probabilities for X.
    The predicted class probabilities of an input sample are computed as
    the mean predicted class probabilities of the trees in the forest. The
    class probability of a single tree is the fraction of samples of the same
    class in a leaf.
    Parameters
    ----------
    X : array-like or sparse matrix of shape = [n_samples, n_features]
        The input samples. Internally, its dtype will be converted to
        ``dtype=np.float32``. If a sparse matrix is provided, it will be
        converted into a sparse ``csr_matrix``.
    Returns
    -------
    p : array of shape = [n_samples, n_classes], or a list of n_outputs
        such arrays if n_outputs > 1.
        The class probabilities of the input samples. The order of the
        classes corresponds to that in the attribute `classes_`.
    """
    check_is_fitted(self, "estimators_")

    # Check data
    if self.check_input:
        X = self._validate_X_predict(X)

    # Assign chunk of trees to jobs
    n_jobs, _, _ = _partition_estimators(self.n_estimators, self.n_jobs)

    all_proba = Parallel(n_jobs=n_jobs, verbose=self.verbose)(
        delayed(e.predict_proba)(X) for e in self.estimators_
    )

    all_proba = np.sum(all_proba, axis=0) / len(self.estimators_)

    return all_proba