BNN weights collapse

I'm working on a BNN regression problem, but the output seems to be collapsing to just learning the mean without any randomness. I don't think it's a problem with the data/model architecture, because I don't have any problems with the same model written with Theano and PyMC3.

I'll insert the model specification code below in case I've made a dumb mistake anywhere. I'd appreciate any insight into what could be wrong. I'd love to be able to switch to TFP; it seems much faster and more powerful!

Edit: Sorry about the weird indentation - I just copied and pasted and bleh. But trust that isn't the issue ;)

input = tf.keras.Input(shape=(x.shape[1],))
                h = input
                layers = [h]
                for i in range(n_hidden):
                    h = tfp.layers.DenseFlipout(hidden_size // np.power(2, i), activation=tf.nn.relu,
                                                name='layer_{}'.format(i),
                                                kernel_posterior_fn=default_mean_field_normal_fn(
                                                    loc_initializer=tf.random_normal_initializer(stddev=1.0),
                                                ))(h)
                    layers.append(h)
                    h = tf.keras.layers.Concatenate()(layers)

                loc_output = tfp.layers.DenseFlipout(1, name='loc_output',
                                                     kernel_posterior_fn=default_mean_field_normal_fn(
                                                         loc_initializer=tf.random_normal_initializer(stddev=1.0),
                                                     ))(h)
                model = tf.keras.Model(inputs=input, outputs=loc_output)
                model.summary()

                _loc = model(features)
                labels_distribution = tfp.distributions.Normal(loc=_loc, scale=1.)

                label_probs = labels_distribution.sample(name='label_probs')

                neg_log_likelihood = -tf.reduce_mean(labels_distribution.log_prob(labels))
                kl = sum(model.losses) / len(y)
                elbo_loss = neg_log_likelihood + kl
                elbo_loss = tf.identity(elbo_loss, name='elbo_loss')
                predictions = tf.reduce_mean(_loc, axis=1, name='prediction')
                
                accuracy, accuracy_update_op = tf.metrics.mean_absolute_error(labels=labels,
                                                                              predictions=predictions,
                                                                              name='accuracy')

                optimizer = tf.train.AdamOptimizer(learning_rate)
                train_op = optimizer.minimize(elbo_loss, name='train_op')