def outcome_learner_predict(ml_model_fit, xdata):
    """Internal function to take a fitted custom_model for the outcome nuisance model and generate the predictions.

    Parameters
    ----------
    ml_model_fit :
        Fitted user-specified model
    xdata : array
        Covariate data to generate the predictions with.

    Returns
    -------
    array
        Predicted values for the outcome (probability if binary, and expected value otherwise)
    """
    if hasattr(ml_model_fit, 'predict_proba'):
        g = ml_model_fit.predict_proba(xdata)
        if g.ndim == 1:  # allows support for pygam.LogisticGAM
            return g
        else:
            return g[:, 1]
    elif hasattr(ml_model_fit, 'predict'):
        return ml_model_fit.predict(xdata)
    else:
        raise ValueError("Currently custom_model must have 'predict' or 'predict_proba' attribute")