diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
index 8c33aacd9ff890..c2daf1da6db30c 100644
--- a/src/transformers/modeling_tf_utils.py
+++ b/src/transformers/modeling_tf_utils.py
@@ -1147,6 +1147,7 @@ def __init__(self, config, *inputs, **kwargs):
         self.config = config
         self.name_or_path = config.name_or_path
         self.generation_config = GenerationConfig.from_model_config(config) if self.can_generate() else None
+        self._set_save_spec(self.input_signature)
 
     def get_config(self):
         return self.config.to_dict()
diff --git a/tests/models/auto/test_modeling_tf_auto.py b/tests/models/auto/test_modeling_tf_auto.py
index 537d48a57e48e5..9c284a78aee56c 100644
--- a/tests/models/auto/test_modeling_tf_auto.py
+++ b/tests/models/auto/test_modeling_tf_auto.py
@@ -211,7 +211,7 @@ def test_from_pretrained_with_tuple_values(self):
         config = copy.deepcopy(model.config)
         config.architectures = ["FunnelBaseModel"]
         model = TFAutoModel.from_config(config)
-        model.build()
+        model.build_in_name_scope()
 
         self.assertIsInstance(model, TFFunnelBaseModel)
 
@@ -249,7 +249,7 @@ def test_new_model_registration(self):
                     config = NewModelConfig(**tiny_config.to_dict())
 
                     model = auto_class.from_config(config)
-                    model.build()
+                    model.build_in_name_scope()
 
                     self.assertIsInstance(model, TFNewModel)
 
diff --git a/tests/models/gpt2/test_modeling_tf_gpt2.py b/tests/models/gpt2/test_modeling_tf_gpt2.py
index a88435acba3e63..d636097dc28622 100644
--- a/tests/models/gpt2/test_modeling_tf_gpt2.py
+++ b/tests/models/gpt2/test_modeling_tf_gpt2.py
@@ -445,7 +445,7 @@ def test_onnx_runtime_optimize(self):
                 continue
 
             model = model_class(config)
-            model.build()
+            model.build_in_name_scope()
 
             onnx_model_proto, _ = tf2onnx.convert.from_keras(model, opset=self.onnx_min_opset)
 
diff --git a/tests/models/whisper/test_modeling_tf_whisper.py b/tests/models/whisper/test_modeling_tf_whisper.py
index e7ac93a5adbe6a..9d1a3299349a94 100644
--- a/tests/models/whisper/test_modeling_tf_whisper.py
+++ b/tests/models/whisper/test_modeling_tf_whisper.py
@@ -312,7 +312,7 @@ def test_encoder_sinusoidal_embed_positions(self):
         config = self.model_tester.get_config()
         for model_class in self.all_model_classes:
             model = model_class(config)
-            model.build()
+            model.build_in_name_scope()
 
             embeds = model.get_encoder().embed_positions.get_weights()[0]
             sinusoids = sinusoidal_embedding_init(embeds.shape).numpy()
diff --git a/tests/utils/test_modeling_tf_core.py b/tests/utils/test_modeling_tf_core.py
index ebd5dfda6bc019..7ec2198dd1523f 100644
--- a/tests/utils/test_modeling_tf_core.py
+++ b/tests/utils/test_modeling_tf_core.py
@@ -217,7 +217,7 @@ def test_saved_model_creation_extended(self):
         for model_class in self.all_model_classes[:2]:
             class_inputs_dict = self._prepare_for_class(inputs_dict, model_class)
             model = model_class(config)
-            model.build()
+            model.build_in_name_scope()
             num_out = len(model(class_inputs_dict))
 
             for key in list(class_inputs_dict.keys()):