Adding probability layer to tensorflow model

I'm predicting a multidimensional output of binary variables. I already have a tensorflow model that performs extremely well on the dataset. However, I want confidence estimations on my predictions. I can see in some of the tutorials that you can add a tensorflow_probability dense layer to the end of the model. With that in mind, here is my relevant code.

class MyModel(Model):
def init(self):
super(MyModel, self).init()
self.d1 = Reshape((322, 1))
self.d2 = GRU(16)
self.d3 = Dense(126, activation='relu')

self.c1 = Reshape((322, 1))
self.c2 = GRU(8, return_sequences=True)
self.c3 = GRU(8)
self.c4 = Dense(126, activation='relu')

self.b1 = Reshape((322, 1))
self.b2 = Conv1D(16, 1, activation='relu')
self.b3 = GRU(16)
self.b4 = Dense(126, activation='relu')    


self.a1 = Dense(1024, activation = 'relu')
self.a3 = Dense(126, activation='relu')    


self.out = tfp.layers.DenseFlipout(126, activation="sigmoid")

def call(self, x):

d = self.d1(x)
d = self.d2(d)
d = self.d3(d)
#d = self.d4(d)

c = self.c1(x)
c = self.c2(c)
c = self.c3(c)
c = self.c4(c)

b = self.b1(x)
b = self.b2(b)
b = self.b3(b)
b = self.b4(b)

a = self.a1(x)
a = self.a3(a)

w = tf.keras.layers.Concatenate()([a, b, c, d])
w = tf.keras.layers.GaussianNoise(0.3)(w)
output = self.out(w)
return output

Followed by, to train,

model = MyModel()

model.compile(optimizer=tf.keras.optimizers.Adam(0.0001),
loss=tf.keras.losses.BinaryCrossentropy(),
metrics=['mae'])

model.fit(X_train, y_train, epochs=20, batch_size=256, validation_data = (X_test, y_test))

Originally, self.out was a simple dense sigmoid output layer. I just replaced it with tfp.layers.DenseFlipout. Is this right? It's hard to tell from all the custom loss functions used in the tutorials.

Start with the MyModel definition, especially the tfp.layers.DenseFlipout output layer, then review the compile and fit calls alongside the TensorFlow Probability tutorials mentioned in the issue. Determine whether the layer and loss setup provide the requested confidence estimates, and document the correct training and evaluation criteria.