TypeError: Can not convert a _TensorCoercible into a Tensor or Operation.

When doing the model prediction on test data,

inputs = Input(shape=(X_train.shape[1],))  
rate = Dense(1, activation=tf.exp)(inputs) 
p_y = tfp.layers.DistributionLambda(tfd.Poisson)(rate) 
model_p = Model(inputs=inputs, outputs=p_y) 

def NLL(y_true, y_hat): 
    return -y_hat.log_prob(y_true)

model_p.compile(Adam(learning_rate=0.001), loss=NLL)
model_p.summary()

hist_p = model_p.fit(x=X_train, y=y_train, validation_data=(X_test, y_test), epochs=10, verbose=1)
y_hat_test = model.predict(X_test)

I am hitting the following error:

Can not convert a _TensorCoercible into a Tensor or Operation.
Traceback (most recent call last):
  File "/tmp/2964910836768450552", line 221, in execute
    exec(code, global_dict)
  File "<livy-input-316cf3fa-854b-48d2-b316-a8203b1ea8d6>", line 1, in <module>
    y_hat_test = model.predict(X_test)
  File "/usr/lib/python3.6/site-packages/keras/engine/training_v1.py", line 979, in predict
    use_multiprocessing=use_multiprocessing)
  File "/usr/lib/python3.6/site-packages/keras/engine/training_arrays_v1.py", line 705, in predict
    callbacks=callbacks)
  File "/usr/lib/python3.6/site-packages/keras/engine/training_arrays_v1.py", line 177, in model_iteration
    f = _make_execution_function(model, mode)
  File "/usr/lib/python3.6/site-packages/keras/engine/training_arrays_v1.py", line 547, in _make_execution_function
    return model._make_execution_function(mode)
  File "/usr/lib/python3.6/site-packages/keras/engine/training_v1.py", line 2085, in _make_execution_function
    self._make_predict_function()
  File "/usr/lib/python3.6/site-packages/keras/engine/training_v1.py", line 2075, in _make_predict_function
    **kwargs)
  File "/usr/lib/python3.6/site-packages/keras/backend.py", line 4093, in function
    inputs, outputs, updates=updates, name=name, **kwargs)
  File "/usr/lib/python3.6/site-packages/keras/backend.py", line 3885, in __init__
    with tf.control_dependencies([self.outputs[0]]):
  File "/usr/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 5394, in control_dependencies
    return get_default_graph().control_dependencies(control_inputs)
  File "/usr/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 4848, in control_dependencies
    c = self.as_graph_element(c)
  File "/usr/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3759, in as_graph_element
    return self._as_graph_element_locked(obj, allow_tensor, allow_operation)
  File "/usr/lib/python3.6/site-packages/tensorflow/python/framework/ops.py", line 3848, in _as_graph_element_locked
    (type(obj).__name__, types_str))
TypeError: Can not convert a _TensorCoercible into a Tensor or Operation.

tensorflow version: 2.6.2
tensorflow_probablity version: 0.14.1

Start by reproducing the supplied model and prediction call with TensorFlow 2.6.2 and TensorFlow Probability 0.14.1. Inspect the traceback around Keras prediction construction and the DistributionLambda output; done means identifying a compatible prediction path or a confirmed fix for the TensorCoercible conversion error.