Add ALL_LAYERNORM_LAYERS to match pytorch

huggingface · Nov 21, 2023 · 7e0a351 · 7e0a351
1 parent b42fe29
commit 7e0a351
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/src/transformers/models/idefics/modeling_tf_idefics.py b/src/transformers/models/idefics/modeling_tf_idefics.py
@@ -28,7 +28,7 @@
 from ...modeling_outputs import ModelOutput
 from ...modeling_utils import PretrainedConfig
 from ...modeling_tf_utils import shape_list
-#from ...pytorch_utils import ALL_LAYERNORM_LAYERS
+from ...tf_utils import ALL_LAYERNORM_LAYERS
 from ...utils import (
     add_start_docstrings,
     add_start_docstrings_to_model_forward,
@@ -487,7 +487,7 @@ def call(self, hidden_states):
         return self.weight * hidden_states
 
 
-#ALL_LAYERNORM_LAYERS.append(TFIdeficsRMSNorm)
+ALL_LAYERNORM_LAYERS.append(TFIdeficsRMSNorm)
 
 
 class TFIdeficsEmbedding(tf.keras.layers.Layer):

diff --git a/src/transformers/tf_utils.py b/src/transformers/tf_utils.py
@@ -19,6 +19,7 @@
 
 from .utils import logging
 
+ALL_LAYERNORM_LAYERS = [tf.keras.layers.LayerNormalization]
 
 logger = logging.get_logger(__name__)
Original file line number	Diff line number	Diff line change
Expand Up		@@ -19,6 +19,7 @@

		from .utils import logging

		ALL_LAYERNORM_LAYERS = [tf.keras.layers.LayerNormalization]

		logger = logging.get_logger(__name__)

Expand Down