diff --git a/src/transformers/models/nllb_moe/modeling_nllb_moe.py b/src/transformers/models/nllb_moe/modeling_nllb_moe.py index 567b03c1c4bbe0..b6ea574469ee80 100644 --- a/src/transformers/models/nllb_moe/modeling_nllb_moe.py +++ b/src/transformers/models/nllb_moe/modeling_nllb_moe.py @@ -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"""