Correctly build models and import call_context for older TF versions (#…

…24138)
huggingface · Jun 9, 2023 · 50ed793 · 50ed793
1 parent fe861e5
commit 50ed793
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/transformers/modeling_tf_utils.py b/src/transformers/modeling_tf_utils.py
@@ -81,7 +81,7 @@
     from keras.engine.keras_tensor import KerasTensor
 else:
     from tensorflow.python.keras import backend as K
-    from tensorflow.python.keras.engine import call_context
+    from tensorflow.python.keras.engine.base_layer_utils import call_context
     from tensorflow.python.keras.engine.keras_tensor import KerasTensor
 
 
@@ -1156,8 +1156,8 @@ def build(self, input_shape=None):
         if self.built or call_context().in_call:
             self.built = True
         else:
-            self(self.dummy_inputs, training=False)
             self.built = True
+            self(self.dummy_inputs, training=False)
 
     def __init__(self, config, *inputs, **kwargs):
         super().__init__(*inputs, **kwargs)