How to use the causalml.propensity.ElasticNetPropensityModel function in causalml

'y': y_train,
        'X': X_train,
        'w': w_train,
        'tau': tau_train,
        'b': b_train,
        'e': e_train}
    preds_dict_valid['generated_data'] = {
        'y': y_val,
        'X': X_val,
        'w': w_val,
        'tau': tau_val,
        'b': b_val,
        'e': e_val}

    # Predict p_hat because e would not be directly observed in real-life
    p_model = ElasticNetPropensityModel()
    p_hat_train = p_model.fit_predict(X_train, w_train)
    p_hat_val = p_model.fit_predict(X_val, w_val)

    for base_learner, label_l in zip([BaseSRegressor, BaseTRegressor, BaseXRegressor, BaseRRegressor],
                                     ['S', 'T', 'X', 'R']):
        for model, label_m in zip([LinearRegression, XGBRegressor], ['LR', 'XGB']):
            # RLearner will need to fit on the p_hat
            if label_l != 'R':
                learner = base_learner(model())
                # fit the model on training data only
                learner.fit(X=X_train, treatment=w_train, y=y_train)
                try:
                    preds_dict_train['{} Learner ({})'.format(
                        label_l, label_m)] = learner.predict(X=X_train, p=p_hat_train).flatten()
                    preds_dict_valid['{} Learner ({})'.format(
                        label_l, label_m)] = learner.predict(X=X_val, p=p_hat_val).flatten()

parser.add_argument('--feature-cols', nargs='+', default=PROPENSITY_FEATURES,
                        dest='feature_cols')
    parser.add_argument('--caliper', type=float, default=.2)
    parser.add_argument('--replace', default=False, action='store_true')
    parser.add_argument('--ratio', type=int, default=1)

    args = parser.parse_args()

    logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)

    logger.info('Loading data from {}'.format(args.input_file))
    df = pd.read_csv(args.input_file)
    df[args.treatment_col] = df[args.treatment_col].astype(int)
    logger.info('shape: {}\n{}'.format(df.shape, df.head()))

    pm = ElasticNetPropensityModel(random_state=42)
    w = df[args.treatment_col].values
    X = load_data(data=df,
                  features=args.feature_cols,
                  transformations=PROPENSITY_FEATURE_TRANSFORMATIONS)

    logger.info('Scoring with a propensity model: {}'.format(pm))
    df[SCORE_COL] = pm.fit_predict(X, w)

    logger.info('Balance before matching:\n{}'.format(create_table_one(data=df,
                                                                       treatment_col=args.treatment_col,
                                                                       features=MATCHING_COVARIATES)))
    logger.info('Matching based on the propensity score with the nearest neighbor model')
    psm = NearestNeighborMatch(replace=args.replace,
                               ratio=args.ratio,
                               random_state=42)
    matched = psm.match_by_group(data=df,

Args:
        synthetic_data_func (function): synthetic data generation function
        n (int, optional): number of samples
        estimators (dict of object): dict of names and objects of treatment effect estimators

    Returns:
        (dict): dict of the actual and estimates of treatment effects
    """
    y, X, w, tau, b, e = synthetic_data_func(n=n)

    preds_dict = {}
    preds_dict[KEY_ACTUAL] = tau
    preds_dict[KEY_GENERATED_DATA] = {'y': y, 'X': X, 'w': w, 'tau': tau, 'b': b, 'e': e}

    # Predict p_hat because e would not be directly observed in real-life
    p_model = ElasticNetPropensityModel()
    p_hat = p_model.fit_predict(X, w)

    if estimators:
        for name, learner in estimators.items():
            try:
                preds_dict[name] = learner.fit_predict(X=X, p=p_hat, treatment=w, y=y).flatten()
            except TypeError:
                preds_dict[name] = learner.fit_predict(X=X, treatment=w, y=y).flatten()
    else:
        for base_learner, label_l in zip([BaseSRegressor, BaseTRegressor, BaseXRegressor, BaseRRegressor],
                                         ['S', 'T', 'X', 'R']):
            for model, label_m in zip([LinearRegression, XGBRegressor], ['LR', 'XGB']):
                learner = base_learner(model())
                model_name = '{} Learner ({})'.format(label_l, label_m)
                try:
                    preds_dict[model_name] = learner.fit_predict(X=X, p=p_hat, treatment=w, y=y).flatten()