Cast using FunctionType instead of input_signature

PiperOrigin-RevId: 506724481
tensorflow · Feb 6, 2023 · eabec3f · eabec3f
1 parent d7259ca
commit eabec3f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/...rflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder_test.py b/...rflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder_test.py
@@ -172,10 +172,10 @@ def test_input_signature_enforced(self):
         tf.TensorSpec.from_tensor(x))
 
     state = encoder.initial_state()
-    with self.assertRaises(ValueError):
+    with self.assertRaises(TypeError):
       bad_x = tf.stack([x, x])
       encoder.encode(bad_x, state)
-    with self.assertRaises(ValueError):
+    with self.assertRaises(TypeError):
       bad_state = state + (x,)
       encoder.encode(x, bad_state)
     encoded_x = encoder.encode(x, state)