ufunc 'isnan' not supported for the input types in DML "effect()" function

Thank you for the package and such huge effort. I am trying to do below estimation,

variables are,
features: ['X1', 'X2', 'X3', 'X4', 'X5'], output: ['Y'], treatment: ['T_1', 'T_2']
Here, Type is categorical and values are (0, 1, 2)

test_seg = dml_test_X.iloc[[2,4]] # third and fifth rows
print(test_seg)
dml_est.effect(test_seg, T0=0, T1=1)

   X1      X2      X3           X4       X5

6 27 1 77.99 4.193 131.126667
10 60 1 76.65 3.717 223.173417

X1- continuous
X2- categorical
X3- continuous
X4- continuous
X5- continuous

treatment- categorical
Y- continuous

ERROR:
TypeError: ufunc 'isnan' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''

Possible reason:
X2 and Treatment are categorical (pandas - dtype('o')) and np.isnan() throws error for "category" data type.

Possible Solution:
replace np.isnan with pd.isna, which supports category dtypes?

No source file or test is named in the report. First reproduce dml_est.effect(test_seg, T0=0, T1=1) with categorical X2 and treatment values, then trace the np.isnan call that rejects the pandas categorical data. Done means the effect call accepts these inputs without the TypeError and coverage verifies the behavior.