Fix broken TFMO tests after Sequential serialization fix.

PiperOrigin-RevId: 544483744

Author: nkovela1

tensorflow_model_optimization/python/core/quantization/keras/graph_transformations/model_transformer_test.py

@@ -463,12 +463,13 @@ def replacement(self, match_layer):
       inp = keras.layers.Input((3,))
       out = keras.layers.Dense(2, activation='relu')(inp)
       model_fused = keras.Model(inp, out)
-      # Ensures old Keras serialization format
-      model_fused.use_legacy_config = True
     else:
       model_fused = keras.Sequential(
           [keras.layers.Dense(2, activation='relu', input_shape=(3,))])
 
+    # Ensures old Keras serialization format
+    model_fused.use_legacy_config = True
+
     if model_type == 'functional':
       inp = keras.layers.Input((3,))
       x = keras.layers.Dense(2)(inp)

tensorflow_model_optimization/python/core/quantization/keras/quantize_integration_test.py

@@ -79,10 +79,12 @@ def _remove_keys(config):
             if isinstance(item, dict):
               _remove_keys(item)
 
+    # Ensure the same config format
+    model1.use_legacy_config, model2.use_legacy_config = True, True
     model1_config = model1.get_config()
     model2_config = model2.get_config()
     exclude_keys = exclude_keys or []
-    exclude_keys += ['build_input_shape', 'build_config']
+    exclude_keys += ['build_config']  # Exclude model build information
     _remove_keys(model1_config)
     _remove_keys(model2_config)
     self.assertEqual(model1_config, model2_config)