Tensorflow Probability: saving and loading model

[shifdz]

I am trying to fit a model with TensorFlow probability, for example:

input = Input(shape=(32,32,32,3))
x = tfp.layers.Convolution3DReparameterization(
        64, kernel_size=5, padding='SAME', activation=tf.nn.relu,
        data_format = 'channels_first')(input)
x = tf.keras.layers.MaxPooling3D(pool_size=(2, 2, 2),
                                 strides=(2, 2, 2),
                                 padding='SAME')(x)
x = tf.keras.layers.Flatten()(x)
output = tfp.layers.DenseFlipout(10)(x)

model3 = Model(input, output)
model3.save('tf_test_model3.h5')

When I load the model as model3 = load_model('tf_test_model3.h5'), I get the following error:

ValueError: Unknown layer: Conv3DReparameterization. Please ensure this object is passed to the `custom_objects` argument. 

When I pass it to custom_objects as:

custom_objects= {'Conv3DReparameterization': tfp.layers.Convolution3DReparameterization}
model3 = load_model('tf_test_model3.h5', custom_objects=custom_objects)

I get the following error: TypeError: 'str' object is not callable

How can I fix this? Please note that I want to save the whole architecture. Not just the weights.

[bobosfw]

Hey, I encountered the problem same as you posted, How do you fixed the problem? Thank you