add no split modules for xlmrobertaxl (#31223)

huggingface · Jun 4, 2024 · 2965b20 · 2965b20
1 parent 821b772
commit 2965b20
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py b/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py
@@ -572,6 +572,7 @@ class XLMRobertaXLPreTrainedModel(PreTrainedModel):
 
     config_class = XLMRobertaXLConfig
     base_model_prefix = "roberta"
+    _no_split_modules = ["XLMRobertaXLEmbeddings", "XLMRobertaXLSelfAttention"]
 
     # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
     def _init_weights(self, module):