Added input_shape param when testing layers

tensorflow · Aug 8, 2021 · 63501a0 · 63501a0
1 parent 75dae2d
commit 63501a0
Showing 1 changed file with 2 additions and 1 deletion.
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.as_list()[1:]),
         layer_class(filters=2, kernel_size=1, data_format=self.data_format)])
 
     net.compile(loss='mse', optimizer='adam')