Merge pull request #4 from togethercomputer/fix_rotar_emb

Replace RotaryEmbedding with GPTNeoXRotaryEmbedding
togethercomputer · Aug 10, 2023 · d0027d5 · d0027d5
2 parents 43a7562 + a6553fe
commit d0027d5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -114,7 +114,7 @@ def __init__(self, config):
         self._init_bias(config.max_position_embeddings)
 
         self.register_buffer("masked_bias", torch.tensor(-1e9), persistent=False)
-        self.rotary_emb = RotaryEmbedding(
+        self.rotary_emb = GPTNeoXRotaryEmbedding(
             self.rotary_ndims, config.max_position_embeddings, base=config.rotary_emb_base, device = 'cuda:0'
         )
         self.register_buffer(