def _validate_output_feature_name_from_train_stats(output_feature_name,
    train_stats_per_model):
    """Validate prediction output_feature_name from model train stats and return it as list.

    :param output_feature_name: output_feature_name containing ground truth
    :param train_stats_per_model: list of per model train stats
    :return output_feature_names: list of output_feature_name(s) containing ground truth
    """
    output_feature_names_set = set()
    for ls in train_stats_per_model:
        for _, values in ls.items():
            for key in values:
                output_feature_names_set.add(key)
    try:
        if output_feature_name in output_feature_names_set:
            return [output_feature_name]
        else:
            return output_feature_names_set
    except TypeError:
        return output_feature_names_set
