Skip to content

Commit

Permalink
Update modeling_gpt_neox.py (#17575)
Browse files Browse the repository at this point in the history
I'm guessing that the intention was to have the `_no_split_modules` class attribute for `GPTNeoXPreTrainedModel` to be set to `["GPTNeoXLayer"]`, akin to how its set as `["GPTJBlock"]` for `GPTJPreTrainedModel`.

If this is incorrect, please feel free to just close the PR.

Thanks!
  • Loading branch information
willfrey committed Jun 13, 2022
1 parent a1344db commit 5483388
Showing 1 changed file with 1 addition and 0 deletions.
1 change: 1 addition & 0 deletions src/transformers/models/gpt_neox/modeling_gpt_neox.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@ class GPTNeoXPreTrainedModel(PreTrainedModel):
config_class = GPTNeoXConfig
base_model_prefix = "gpt_neox"
supports_gradient_checkpointing = True
_no_split_modules = ["GPTNeoXLayer"]

def _init_weights(self, module):
"""Initialize the weights"""
Expand Down

0 comments on commit 5483388

Please sign in to comment.