Saving and Loading Bayesian Neural Network

I'd been looking at the Tensorflow Probability library and trying to modify the example in bayesian network example, hoping that I can save checkpoints and then restore them.

I first started trying to use tf.train.Checkpoint but, although, I was not getting any error when saving nor when restoring, it didnt seem to restart the training from the previous checkpoint as the accuracy was completely different value. I then tried using tf.keras.models.model.save, which again does save a file, but when trying to restore, I get the error: ValueError: Unknown layer: Conv2DFlipout when it is trying to deserialise the layer.

I then proceeded to add the custom objects Conv2DFlipout and DenseFlipout, and now I am getting the following error, and I am now a bit lost :( on what to do

File "C:/Users/ezzgm/Documents/BayesianNeural Network/bayesian_NN_vKerasSaver.py", line 188, in main
model = tf.keras.models.load_model(FLAGS.model_dir+'checkpoint.hdf5')

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\engine\saving.py", line 230, in load_model
model = model_from_config(model_config, custom_objects=custom_objects)

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\engine\saving.py", line 310, in model_from_config
return deserialize(config, custom_objects=custom_objects)

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\layers\serialization.py", line 64, in deserialize
printable_module_name='layer')

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\utils\generic_utils.py", line 173, in deserialize_keras_object
list(custom_objects.items())))

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\engine\network.py", line 1302, in from_config
process_node(layer, node_data)

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\engine\network.py", line 1260, in process_node
layer(input_tensors[0], **kwargs)

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow\python\keras\engine\base_layer.py", line 746, in call
self.build(input_shapes)

File "C:\Users\ezzgm\AppData\Local\conda\conda\envs\tensorflow1.12\lib\site-packages\tensorflow_probability\python\layers\conv_variational.py", line 188, in build
self.trainable, self.add_variable)

File "C:/Users/ezzgm/AppData/Local/conda/conda/envs/tensorflow1.12/lib/site-packages/tensorflow_probability/python/layers/util.py", line 186, in _fn
loc, scale = loc_scale_fn(dtype, shape, name, trainable, add_variable_fn)

TypeError: 'str' object is not callable

This is the modified code of the Bayesian Neural Network (I am trying to save in the first 10 steps just for debugin purposes and then rerunning the code to try and load the saved model). Any help on this will be very much appreciated!

`

   if FLAGS.architecture == "resnet":
       model_fn = bayesian_resnet.bayesian_resnet
   else:
       model_fn = bayesian_vgg.bayesian_vgg


  model = model_fn(
       IMAGE_SHAPE,
       num_classes=4,
       kernel_posterior_scale_mean=FLAGS.kernel_posterior_scale_mean,
       kernel_posterior_scale_constraint=FLAGS.kernel_posterior_scale_constraint)


  #check if saved checkpoint exists
  exists = os.path.isfile(FLAGS.model_dir+"checkpoint.hdf5")
  if exists:
       with   tf.keras.utils.CustomObjectScope({'DenseFlipout':tfp.layers.DenseFlipout,'Conv2DFlipout':tfp.layers.Convolution2DFlipout}):
      model = tf.keras.models.load_model(FLAGS.model_dir+'checkpoint.hdf5') 

  logits = model(images)
  labels_distribution = tfd.Categorical(logits=logits`

Start in bayesian_NN_vKerasSaver.py at main, then inspect the bayesian_resnet.bayesian_resnet and bayesian_vgg.bayesian_vgg entry points used to construct the model. Reproduce the save and load paths with the TensorFlow 1.12 environment and compare checkpoint restoration with HDF5 deserialization. Done means the issue has a documented, reproducible save-and-restore path or a confirmed minimal failure case.