Fix casting-dependent tests

PiperOrigin-RevId: 507607482

Author: faizan-m

tensorflow_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, ValueError)):
       bad_x = tf.stack([x, x])
       encoder.encode(bad_x, state)
-    with self.assertRaises(ValueError):
+    with self.assertRaises((TypeError, ValueError)):
       bad_state = state + (x,)
       encoder.encode(x, bad_state)
     encoded_x = encoder.encode(x, state)