[`Nllb-Moe`] Fix nllb moe accelerate issue (#23758)
fix nllb moe accelerate issue
This commit is contained in:
parent
d685e330b5
commit
f67dac97bd
|
@ -856,7 +856,7 @@ class NllbMoePreTrainedModel(PreTrainedModel):
|
|||
config_class = NllbMoeConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["NllbMoeAttention"]
|
||||
_no_split_modules = ["NllbMoeEncoderLayer", "NllbMoeDecoderLayer"]
|
||||
|
||||
def _init_weights(self, module):
|
||||
"""Initialize the weights"""
|
||||
|
|
Loading…
Reference in New Issue