model.losses is not accessible with a model built via tf.keras.Model(inputs, outputs)

I am using tensorflow==2.4.0, tensorflow-estimator==2.4.0, tensorflow-probability==0.12.2.

The first two versions of this toy example do work, but the third does not. I believe the third should work too. Instead it gives me

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: conv2d_flipout_4/divergence_kernel:0

The three examples are:

import tensorflow as tf
import tensorflow_probability as tfp

# This first version works
class MyModel(tf.keras.Model):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = tf.keras.layers.Reshape([32, 32, 3])
        self.layer2 = tfp.layers.Convolution2DFlipout(64, kernel_size=5, padding="SAME", activation=tf.nn.relu)
        self.layer3 = tf.keras.layers.MaxPooling2D(pool_size=[2, 2], strides=[2, 2], padding="SAME")
        self.layer4 = tf.keras.layers.Flatten()
        self.layer5 = tfp.layers.DenseFlipout(10)

    def call(self, inputs):
        x = self.layer1(inputs)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)
        return x


model = MyModel()
kl = sum(model.losses)
print("model.losses worked this time.")


# This second version works
model = tf.keras.Sequential(
    [
        tf.keras.layers.Reshape([32, 32, 3]),
        tfp.layers.Convolution2DFlipout(64, kernel_size=5, padding="SAME", activation=tf.nn.relu),
        tf.keras.layers.MaxPooling2D(pool_size=[2, 2], strides=[2, 2], padding="SAME"),
        tf.keras.layers.Flatten(),
        tfp.layers.DenseFlipout(10),
    ]
)
kl = sum(model.losses)
print("model.losses worked this time, too.")


# This third version fails
inp = tf.keras.Input(shape=(32, 32, 3))
i = inp
i = tf.keras.layers.Reshape([32, 32, 3])(i)
i = tfp.layers.Convolution2DFlipout(64, kernel_size=5, padding="SAME", activation=tf.nn.relu)(i)
i = tf.keras.layers.MaxPooling2D(pool_size=[2, 2], strides=[2, 2], padding="SAME")(i)
i = tf.keras.layers.Flatten()(i)
i = tfp.layers.DenseFlipout(10)(i)
out = i
model = tf.keras.Model(inputs=inp, outputs=out)
# The following line fails
kl = sum(model.losses)
print("This never prints.")

Start by running the three model-construction examples in the issue with TensorFlow 2.4.0, TensorFlow Estimator 2.4.0, and TensorFlow Probability 0.12.2, focusing on the failing sum(model.losses) call in the functional-model example. Trace how model.losses handles the Flipout layers and Graph tensors, then verify that the functional example works without breaking the subclassed and Sequential examples.