Problem defining model object in bayesian_vgg.py

Hello all,

my following code

self._input = tf.placeholder(tf.float32, shape=[None, 28, 28, 1])
self._training = tf.placeholder_with_default(False, shape=[])
self._labels = tf.placeholder(tf.float32, shape=[None, 10])

out = self._input
   
out = tfp.layers.Convolution2DFlipout(6,kernel_size=5,padding="SAME",activation=tf.nn.relu)(out)
out = tf.layers.max_pooling2d(out, pool_size=2, strides=2) # 12
out = tfp.layers.Convolution2DFlipout(16,kernel_size=5,padding="SAME",activation=tf.nn.relu)(out)
out = tf.layers.max_pooling2d(out, pool_size=2, strides=2) # 4
out = tf.contrib.layers.flatten(out)
out = tfp.layers.DenseFlipout(120, activation=tf.nn.relu)(out)
out = tfp.layers.DenseFlipout(84, activation=tf.nn.relu)(out)
out = tfp.layers.DenseFlipout(10)(out)
    
image = tf.keras.layers.Input(shape=[None, 28, 28, 1], dtype='float32')
model = tf.keras.Model(inputs=image, outputs=out)

returns the following error

Output tensors to a Model must be the output of a TensorFlow Layer (thus holding past layer metadata). Found: Tensor("dense_flipout_1/BiasAdd:0", shape=(?, 10), dtype=float32)

Anyone has an idea on how to solve that?

Thanks so much