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
You need to disable eager execution: tf.compat.v1.disable_eager_execution().