Unverified Commit f67dac97 authored by Younes Belkada's avatar Younes Belkada Committed by GitHub
Browse files

[`Nllb-Moe`] Fix nllb moe accelerate issue (#23758)

fix nllb moe accelerate issue
parent d685e330
...@@ -856,7 +856,7 @@ class NllbMoePreTrainedModel(PreTrainedModel): ...@@ -856,7 +856,7 @@ class NllbMoePreTrainedModel(PreTrainedModel):
config_class = NllbMoeConfig config_class = NllbMoeConfig
base_model_prefix = "model" base_model_prefix = "model"
supports_gradient_checkpointing = True supports_gradient_checkpointing = True
_no_split_modules = ["NllbMoeAttention"] _no_split_modules = ["NllbMoeEncoderLayer", "NllbMoeDecoderLayer"]
def _init_weights(self, module): def _init_weights(self, module):
"""Initialize the weights""" """Initialize the weights"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment