Bootstrap and DeepIV does not work on Multiple Workers

I tried adding the bootstrap wrapper to the file Deep IV Examples.ipynb with n_jobs more than 1, but it gives me some errors:

deepIvEst = DeepIVEstimator(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/ 
                            s1=100, # number of epochs to train treatment model
                            s2=100) # number of epochs to train response model
boot_est = BootstrapEstimator(deepIvEst, n_bootstrap_samples=100, n_jobs=4)
boot_est.fit(Y=y,T=t,X=x,Z=z)

[Parallel(n_jobs=4)]: Using backend ThreadingBackend with 4 concurrent workers.
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-14-2528bf95d3ca> in <module>
----> 1 boot_est.fit(Y=y,T=t,X=x,Z=z)

~/Dropbox (Cambridge University)/R301/EconML/econml/bootstrap.py in fit(self, *args, **named_args)
     56         Parallel(n_jobs=self._n_jobs, prefer='threads', verbose=3)(
     57             (obj.fit, [arg[inds] for arg in args], {arg: named_args[arg][inds] for arg in named_args})
---> 58             for obj, inds in zip(self._instances, indices)
     59         )
     60         return self

~/anaconda3/lib/python3.6/site-packages/joblib-0.13.2-py3.6.egg/joblib/parallel.py in __call__(self, iterable)
    932 
    933             with self._backend.retrieval_context():
--> 934                 self.retrieve()
    935             # Make sure that we get a last message telling us we are done
    936             elapsed_time = time.time() - self._start_time

~/anaconda3/lib/python3.6/site-packages/joblib-0.13.2-py3.6.egg/joblib/parallel.py in retrieve(self)
    831             try:
    832                 if getattr(self._backend, 'supports_timeout', False):
--> 833                     self._output.extend(job.get(timeout=self.timeout))
    834                 else:
    835                     self._output.extend(job.get())

~/anaconda3/lib/python3.6/multiprocessing/pool.py in get(self, timeout)
    642             return self._value
    643         else:
--> 644             raise self._value
    645 
    646     def _set(self, i, obj):

~/anaconda3/lib/python3.6/multiprocessing/pool.py in worker(inqueue, outqueue, initializer, initargs, maxtasks, wrap_exception)
    117         job, i, func, args, kwds = task
    118         try:
--> 119             result = (True, func(*args, **kwds))
    120         except Exception as e:
    121             if wrap_exception and func is not _helper_reraises_exception:

~/anaconda3/lib/python3.6/site-packages/joblib-0.13.2-py3.6.egg/joblib/_parallel_backends.py in __call__(self, *args, **kwargs)
    565     def __call__(self, *args, **kwargs):
    566         try:
--> 567             return self.func(*args, **kwargs)
    568         except KeyboardInterrupt:
    569             # We capture the KeyboardInterrupt and reraise it as

~/anaconda3/lib/python3.6/site-packages/joblib-0.13.2-py3.6.egg/joblib/parallel.py in __call__(self)
    223         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    224             return [func(*args, **kwargs)
--> 225                     for func, args, kwargs in self.items]
    226 
    227     def __len__(self):

~/anaconda3/lib/python3.6/site-packages/joblib-0.13.2-py3.6.egg/joblib/parallel.py in <listcomp>(.0)
    223         with parallel_backend(self._backend, n_jobs=self._n_jobs):
    224             return [func(*args, **kwargs)
--> 225                     for func, args, kwargs in self.items]
    226 
    227     def __len__(self):

~/Dropbox (Cambridge University)/R301/EconML/econml/deepiv.py in fit(self, Y, T, X, Z)
    313         n_components = self._n_components
    314 
--> 315         treatment_network = self._m(z_in, x_in)
    316 
    317         # the dimensionality of the output of the network

<ipython-input-10-df8c54d69989> in <lambda>(z, x)
      1 deepIvEst = DeepIVEstimator(n_components = 10, # number of gaussians in our mixture density network
----> 2                             m = lambda z, x : treatment_model(keras.layers.concatenate([z,x])), # treatment model
      3                             h = lambda t, x : response_model(keras.layers.concatenate([t,x])),  # response model
      4                             n_samples = 1, # number of samples to use to estimate the response
      5                             use_upper_bound_loss = False, # whether to use an approximation to the true loss

~/anaconda3/lib/python3.6/site-packages/keras/engine/base_layer.py in __call__(self, inputs, **kwargs)
    455             # Actually call the layer,
    456             # collecting output(s), mask(s), and shape(s).
--> 457             output = self.call(inputs, **kwargs)
    458             output_mask = self.compute_mask(inputs, previous_mask)
    459 

~/anaconda3/lib/python3.6/site-packages/keras/engine/network.py in call(self, inputs, mask)
    562             return self._output_tensor_cache[cache_key]
    563         else:
--> 564             output_tensors, _, _ = self.run_internal_graph(inputs, masks)
    565             return output_tensors
    566 

~/anaconda3/lib/python3.6/site-packages/keras/engine/network.py in run_internal_graph(self, inputs, masks)
    719                                     kwargs['mask'] = computed_mask
    720                             output_tensors = to_list(
--> 721                                 layer.call(computed_tensor, **kwargs))
    722                             output_masks = layer.compute_mask(computed_tensor,
    723                                                               computed_mask)

~/anaconda3/lib/python3.6/site-packages/keras/layers/core.py in call(self, inputs)
    877 
    878     def call(self, inputs):
--> 879         output = K.dot(inputs, self.kernel)
    880         if self.use_bias:
    881             output = K.bias_add(output, self.bias, data_format='channels_last')

~/anaconda3/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py in dot(x, y)
   1083         out = tf.sparse_tensor_dense_matmul(x, y)
   1084     else:
-> 1085         out = tf.matmul(x, y)
   1086     return out
   1087 

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/ops/math_ops.py in matmul(a, b, transpose_a, transpose_b, adjoint_a, adjoint_b, a_is_sparse, b_is_sparse, name)
   2385       are both set to True.
   2386   """
-> 2387   with ops.name_scope(name, "MatMul", [a, b]) as name:
   2388     if transpose_a and adjoint_a:
   2389       raise ValueError("Only one of transpose_a and adjoint_a can be True.")

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in __enter__(self)
   6081       if self._values is None:
   6082         self._values = []
-> 6083       g = _get_graph_from_inputs(self._values)
   6084       self._g_manager = g.as_default()
   6085       self._g_manager.__enter__()

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in _get_graph_from_inputs(op_input_list, graph)
   5711         graph = graph_element.graph
   5712       elif original_graph_element is not None:
-> 5713         _assert_same_graph(original_graph_element, graph_element)
   5714       elif graph_element.graph is not graph:
   5715         raise ValueError("%s is not from the passed-in graph." % graph_element)

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/ops.py in _assert_same_graph(original_item, item)
   5647   if original_item.graph is not item.graph:
   5648     raise ValueError("%s must be from the same graph as %s." % (item,
-> 5649                                                                 original_item))
   5650 
   5651 

ValueError: Tensor("dense_1/kernel:0", shape=(2, 128), dtype=float32_ref) must be from the same graph as Tensor("concatenate_1/concat:0", shape=(?, 2), dtype=float32).

However, the problem is solved when I change n_jobs=1:

boot_est = BootstrapEstimator(deepIvEst, n_bootstrap_samples=10, n_jobs=1)

Is this a currency issue (i.e. DeepIV can only be trained on one job)? But it shouldn't be because in Bootstrapping, you should be able to do parallel processing ...