DRtester does not work for binary treatment AND binary outcome

Hi, here is the code to reproduce the error:

import numpy as np
import pandas as pd
import scipy.stats as st

from econml.metalearners import TLearner
from econml.dml import DML

from econml.validate.drtester import DRtester

from sklearn.ensemble import (RandomForestClassifier, RandomForestRegressor,
                              GradientBoostingClassifier, GradientBoostingRegressor, 
                              AdaBoostClassifier, AdaBoostRegressor)

from sklearn.svm import SVC, SVR
from sklearn.neural_network import MLPClassifier, MLPRegressor
from lightgbm import LGBMClassifier, LGBMRegressor
                                
    # Classifiers
from sklearn.dummy import DummyClassifier

np.random.seed(123)

N = 20000  # number of units
K = 5  # number of covariates
num_treatments = 2 # number of treatments (excluding control)

# Generate random Xs
X_mu = np.zeros(5)  # Means of Xs
# Random covariance matrix of Xs
X_sig = np.diag(np.random.rand(5))
X = st.multivariate_normal(X_mu, X_sig).rvs(N)

# Effect of Xs on outcome
X_beta = np.random.uniform(0, 5, K)
# Effect of treatment on outcomes
D_beta = np.array([0, 1, 2])
# Effect of treatment on outcome conditional on X1
DX1_beta = np.array([0, 0, 3])

# Generate treatments based on X and random noise
beta_treat = np.random.uniform(-1, 1, (num_treatments + 1, K))
D1 = np.zeros((N, num_treatments + 1))
for k in range(num_treatments + 1):
    D1[:, k] = X @ beta_treat[k, :] + np.random.gumbel(0, 1, N)
D = np.array([np.where(D1[i, :] == np.max(D1[i, :]))[0][0] for i in range(N)])
D_dum = pd.get_dummies(D)

# Generate Y (based on X, D, and random noise)
Y_sig = 1  # Variance of random outcome noise
Y = X @ X_beta + (D_dum @ D_beta) + X[:, 1] * (D_dum @ DX1_beta) + np.random.normal(0, Y_sig, N)
Y = Y.to_numpy()

# Split into training/validation samples
train_prop = .5
train_N = np.ceil(train_prop * N)
ind = np.array(range(N))
train_ind = np.random.choice(N, int(train_N), replace=False)
val_ind = ind[~np.isin(ind, train_ind)]

Xtrain, Dtrain, Ytrain = X[train_ind], D[train_ind], Y[train_ind]
Xval, Dval, Yval = X[val_ind], D[val_ind], Y[val_ind]



model_regression = GradientBoostingRegressor(random_state=0)
model_propensity = RandomForestClassifier(random_state=0)

nval = len(Dval)
ntrain = len(Dtrain)
Dval = np.random.choice(2, nval)
Dtrain = np.random.choice(2, ntrain)

nval = len(Yval)
ntrain = len(Ytrain)
Yval = np.random.choice(2, nval)
Ytrain = np.random.choice(2, ntrain)

est_dm = DML(model_y = model_propensity, 
             model_t= model_propensity,
             model_final = model_regression,
             discrete_outcome=True,
             discrete_treatment=True,
             cv=5)

est_dm.fit(Ytrain, Dtrain, X=Xtrain)

# Initialize DRTester and fit/predict nuisance models
dml_tester = DRtester(
    model_regression=model_regression, 
    model_propensity=model_propensity,
    cate=est_dm
).fit_nuisance(Xval, Dval, Yval, Xtrain, Dtrain, Ytrain)

res_dml = dml_tester.evaluate_all(Xval, Xtrain)
res_dml.summary()

Error message:

ValueError                                Traceback (most recent call last)
/tmp/ipykernel_28104/1480639724.py in <module>
     92 ).fit_nuisance(Xval, Dval, Yval, Xtrain, Dtrain, Ytrain)
     93 
---> 94 res_dml = dml_tester.evaluate_all(Xval, Xtrain)
     95 res_dml.summary()

[~/.local/lib/python3.10/site-packages/econml/validate/drtester.py](http://localhost:8888/home/yvdc/.local/lib/python3.10/site-packages/econml/validate/drtester.py) in evaluate_all(self, Xval, Xtrain, n_groups)
    594             self.get_cate_preds(Xval, Xtrain)
    595 
--> 596         blp_res = self.evaluate_blp()
    597         cal_res = self.evaluate_cal(n_groups=n_groups)
    598         qini_res = self.evaluate_uplift(metric='qini')

[~/.local/lib/python3.10/site-packages/econml/validate/drtester.py](http://localhost:8888/home/yvdc/.local/lib/python3.10/site-packages/econml/validate/drtester.py) in evaluate_blp(self, Xval, Xtrain)
    462 
    463         if self.n_treat == 1:  # binary treatment
--> 464             reg = OLS(self.dr_val_, add_constant(self.cate_preds_val_)).fit()
    465             params = [reg.params[1]]
    466             errs = [reg.bse[1]]

[~/.local/lib/python3.10/site-packages/statsmodels/tools/tools.py](http://localhost:8888/home/yvdc/.local/lib/python3.10/site-packages/statsmodels/tools/tools.py) in add_constant(data, prepend, has_constant)
    191         x = x[:, None]
    192     elif x.ndim > 2:
--> 193         raise ValueError('Only implemented for 2-dimensional arrays')
    194 
    195     is_nonzero_const = np.ptp(x, axis=0) == 0

ValueError: Only implemented for 2-dimensional arrays

It does work when the outcome y is continuous.

Reproduce the provided binary-treatment and binary-outcome example, then inspect econml/validate/drtester.py at evaluate_all and evaluate_blp, along with the statsmodels add_constant call in the traceback. Determine why the CATE predictions produce an unsupported shape, and verify that evaluate_all completes for binary outcomes while the existing continuous-outcome behavior remains intact.