GPTNeo: handle padded wte (huggingface#11078)

leogao2 · Apr 6, 2021 · 666b527 · 666b527
1 parent 6ab7d1a
commit 666b527
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/transformers/models/gpt_neo/modeling_gpt_neo.py b/src/transformers/models/gpt_neo/modeling_gpt_neo.py
@@ -112,6 +112,10 @@ def load_tf_weights_in_gpt_neo(model, config, gpt_neo_checkpoint_path):
         if name[-1] == "w" and name[-2] in ["out_proj", "k_proj", "q_proj", "v_proj", "c_proj", "c_fc"]:
             array = array.transpose()
 
+        if name == ["wte"]:
+            # if vocab is padded, then trim off the padding embeddings
+            array = array[:50257]
+
         try:
             assert (
                 pointer.shape == array.shape