Fix PT TF ViTMAE (huggingface#16766)

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
elusenji · Jun 12, 2022 · 59c39cd · 59c39cd
1 parent f5eea6a
commit 59c39cd
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/src/transformers/models/vit_mae/modeling_tf_vit_mae.py b/src/transformers/models/vit_mae/modeling_tf_vit_mae.py
@@ -860,7 +860,9 @@ def __init__(self, config, num_patches, **kwargs):
 
         self.decoder_norm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="decoder_norm")
         self.decoder_pred = tf.keras.layers.Dense(
-            config.patch_size**2 * config.num_channels, name="decoder_pred"
+            config.patch_size**2 * config.num_channels,
+            kernel_initializer=get_initializer(config.initializer_range),
+            name="decoder_pred",
         )  # encoder to decoder
         self.config = config
         self.num_patches = num_patches

diff --git a/src/transformers/models/vit_mae/modeling_vit_mae.py b/src/transformers/models/vit_mae/modeling_vit_mae.py
@@ -756,7 +756,7 @@ def __init__(self, config, num_patches):
             [ViTMAELayer(decoder_config) for _ in range(config.decoder_num_hidden_layers)]
         )
 
-        self.decoder_norm = nn.LayerNorm(config.decoder_hidden_size)
+        self.decoder_norm = nn.LayerNorm(config.decoder_hidden_size, eps=config.layer_norm_eps)
         self.decoder_pred = nn.Linear(
             config.decoder_hidden_size, config.patch_size**2 * config.num_channels, bias=True
         )  # encoder to decoder