DeepIV.fit with Inference='bootstrap' throws error

[km-Poonacha]

I keep running into a TypeError when I include the inference option in the DeepIV model. This does not occur when I do not have the inference option in the .fit method.

Code to replicate the issue (adapted from the DeepIV example) :

import numpy as np
e = np.random.normal(size=(n,))
X = np.random.uniform(low=0.0, high=10.0, size=(n,))
Z = np.random.uniform(low=0.0, high=10.0, size=(n,))
T = np.sqrt((X+2) * Z) + e
Y = T*T / 10 - X*T / 10 + e

import keras
from econml.iv.nnet import DeepIV

treatment_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(64, activation='relu'),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(32, activation='relu'),
                                    keras.layers.Dropout(0.17)])
response_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(64, activation='relu'),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(32, activation='relu'),
                                  keras.layers.Dropout(0.17),
                                  keras.layers.Dense(1)])
keras_fit_options = { "epochs": 30,
                      "validation_split": 0.1,
                      "callbacks": [keras.callbacks.EarlyStopping(patience=2, restore_best_weights=True)]}

deepIvEst = DeepIV(n_components = 10, # number of gaussians in our mixture density network
                   m = lambda Z, X : treatment_model(keras.layers.concatenate([Z,X])), # treatment model
                   h = lambda T, X : response_model(keras.layers.concatenate([T,X])),  # response model
                   n_samples = 1, # number of samples to use to estimate the response
                   use_upper_bound_loss = False, # whether to use an approximation to the true loss
                   n_gradient_samples = 1, # number of samples to use in second estimate of the response (to make loss estimate unbiased)
                   optimizer='adam', # Keras optimizer to use for training - see https://keras.io/optimizers/ 
                   first_stage_options = keras_fit_options, # options for training treatment model
                   second_stage_options = keras_fit_options) # options for training response model

deepIvEst.fit(Y=Y,T=T,X=X,Z=Z, inference="bootstrap")

The error encountered is :

TypeError: Exception encountered when calling layer "lambda_261" (type Lambda).

'dict' object is not callable

Call arguments received by layer "lambda_261" (type Lambda):
• inputs=tf.Tensor(shape=(None, 10), dtype=float32)
• mask=None
• training=None