[FIXED] Tensorflow 2.5.0 - TypeError: An op outside of the function building code is being passed a "Graph" tensor

Issue

I am training a convolutional Bayesian Neural network where I use tfp.layers.Convolution3DFlipout layers. My loss function is as follows:

from tensorflow.keras.losses import binary_crossentropy

def variational_free_energy_loss(model, scale_factor = tf.constant(1.)):
    kl = sum(model.losses) / scale_factor

    def loss(y_true, y_pred):
        bce = binary_crossentropy(y_true, y_pred)
        return bce + kl 
    return loss

I am getting this error:

TypeError: An op outside of the function building code is being passed
a "Graph" tensor. It is possible to have Graph tensors
leak out of the function building context by including a
tf.init_scope in your function building code.
For example, the following function will fail:
  @tf.function
  def has_init_scope():
    my_constant = tf.constant(1.)
    with tf.init_scope():
      added = my_constant * 2
The graph tensor has name: conv3d_flipout_189/divergence_kernel:0

Does anyone know what is causing this error?

tensorflow version: 2.5.0

tensorflow_probability version: 0.13.0

Solution

You need to disable eager execution: tf.compat.v1.disable_eager_execution().

• Check this gist for TFlow 2.4 and TFlow-Prob 0.12.1
• This was also raised as an issue on the tensorflow/probability repo - link

Answered By - pmod