Merge pull request #1386 from Frightera:frighterafix#1384

PiperOrigin-RevId: 390247581
tensorflow · Aug 11, 2021 · 25bc475 · 25bc475
2 parents 0867ae3 + 63501a0
commit 25bc475
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tensorflow_probability/python/layers/conv_variational.py b/tensorflow_probability/python/layers/conv_variational.py
@@ -381,7 +381,7 @@ def _apply_variational_bias(self, inputs):
         # As of Mar 2017, direct addition is significantly slower than
         # bias_add when computing gradients. To use bias_add, we collapse Z
         # and Y into a single dimension to obtain a 4D input tensor.
-        outputs_shape = outputs.shape.as_list()
+        outputs_shape = tf.shape(outputs)
         outputs_4d = tf.reshape(outputs,
                                 [outputs_shape[0], outputs_shape[1],
                                  outputs_shape[2] * outputs_shape[3],

diff --git a/tensorflow_probability/python/layers/conv_variational_test.py b/tensorflow_probability/python/layers/conv_variational_test.py
@@ -609,7 +609,8 @@ def _testLayerInSequential(self, layer_class):  # pylint: disable=invalid-name
     outputs = self.maybe_transpose_tensor(outputs)
 
     net = tf.keras.Sequential([
-        layer_class(filters=2, kernel_size=3, data_format=self.data_format),
+        layer_class(filters=2, kernel_size=3, data_format=self.data_format,
+                    input_shape=inputs.shape[1:]),
         layer_class(filters=2, kernel_size=1, data_format=self.data_format)])
 
     net.compile(loss='mse', optimizer='adam')