title: Tensorflow Probability: saving and loading model
body: <p>I am trying to fit a model with TensorFlow probability, for example:</p>
  <pre><code>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')
  
  </code></pre>
  <p>When I load the model as <code>model3 = load_model('tf_test_model3.h5')</code>, I get the following error:</p>
  <pre><code>ValueError: Unknown layer: Conv3DReparameterization. Please ensure this object is passed to the `custom_objects` argument. 
  </code></pre>
  <p>When I pass it to custom_objects as:</p>
  <pre><code>custom_objects= {'Conv3DReparameterization': tfp.layers.Convolution3DReparameterization}
  model3 = load_model('tf_test_model3.h5', custom_objects=custom_objects)
  </code></pre>
  <p>I get the following error:
  <code>TypeError: 'str' object is not callable</code></p>
  <p>What am I doing wrong? How can I fix this?</p>