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

fix nllb moe accelerate issue
huggingface · May 25, 2023 · f67dac9 · f67dac9
1 parent d685e33
commit f67dac9
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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"""