# Issue with training a MLP with DenseFlipout layers

Hi, i would like for an advice regarding a problem that i am having with a mlp with DenseFlipout layers.
I have been trying to train the neural network for a regression task and i am using denseFlipout layers with mostly the default settings

``````def create_flipout_bnn_model(train_size):
def normal_sp(params):
return tfd.Normal(loc=params[:,0:1], scale=1e-3 + tf.math.softplus(0.05 * params[:,1:2]))

kernel_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / (train_size)
bias_divergence_fn=lambda q, p, _: tfp.distributions.kl_divergence(q, p) / (train_size)

inputs = Input(shape=(1,),name="input layer")

hidden = tfp.layers.DenseFlipout(50,
kernel_divergence_fn=kernel_divergence_fn,
activation="relu",name="DenseFlipout_layer_1")(inputs)
hidden = tfp.layers.DenseFlipout(50,
kernel_divergence_fn=kernel_divergence_fn,
activation="relu",name="DenseFlipout_layer_2")(hidden)
hidden = tfp.layers.DenseFlipout(50,
kernel_divergence_fn=kernel_divergence_fn,
activation="relu",name="DenseFlipout_layer_3")(hidden)
params = tfp.layers.DenseFlipout(2,
kernel_divergence_fn=kernel_divergence_fn,
name="DenseFlipout_layer_5")(hidden)
dist = tfp.layers.DistributionLambda(normal_sp,name = 'normal_sp')(params)

model = Model(inputs=inputs, outputs=dist)

return model

flipout_BNN = create_flipout_bnn_model(train_size=train_size)
{"ft_lang_label":"__label__en","ft_lang_prob":0.65652287,"math_prob":0.9678057,"size":1857,"snap":"2023-14-2023-23","text_gpt3_token_len":499,"char_repetition_ratio":0.1807879,"word_repetition_ratio":0.036809817,"special_character_ratio":0.24340335,"punctuation_ratio":0.22653721,"nsfw_num_words":0,"has_unicode_error":false,"math_prob_llama3":0.9948571,"pos_list":[0,1,2],"im_url_duplicate_count":[null,1,null],"WARC_HEADER":"{\"WARC-Type\":\"response\",\"WARC-Date\":\"2023-04-02T12:05:15Z\",\"WARC-Record-ID\":\"<urn:uuid:48b8a729-3809-4cae-8ea1-858bc7d87e46>\",\"Content-Length\":\"17139\",\"Content-Type\":\"application/http; msgtype=response\",\"WARC-Warcinfo-ID\":\"<urn:uuid:3fbca506-d31f-4f56-820e-2d44fa5589ff>\",\"WARC-Concurrent-To\":\"<urn:uuid:38856515-f4a1-4f66-bb8e-673107792bdf>\",\"WARC-IP-Address\":\"34.107.138.196\",\"WARC-Target-URI\":\"https://discuss.tensorflow.org/t/issue-with-training-a-mlp-with-denseflipout-layers/13077\",\"WARC-Payload-Digest\":\"sha1:L6TWXP4FMLGO5MI55TREILKCVTBYHRBV\",\"WARC-Block-Digest\":\"sha1:E5C4P2S7I76OYXN3YYFIPLDZYY333JDC\",\"WARC-Identified-Payload-Type\":\"text/html\",\"warc_filename\":\"/cc_download/warc_2023/CC-MAIN-2023-14/CC-MAIN-2023-14_segments_1679296950528.96_warc_CC-MAIN-20230402105054-20230402135054-00540.warc.gz\"}"}