fix _no_split_modules for Whisper model (huggingface#22486)

Lucinao540 · Apr 4, 2023 · 48fbd8f · 48fbd8f
1 parent 9006774
commit 48fbd8f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/whisper/modeling_whisper.py b/src/transformers/models/whisper/modeling_whisper.py
@@ -577,7 +577,7 @@ class WhisperPreTrainedModel(PreTrainedModel):
     base_model_prefix = "model"
     main_input_name = "input_features"
     supports_gradient_checkpointing = True
-    _no_split_modules = ["WhisperEncoderLayer"]
+    _no_split_modules = ["WhisperEncoderLayer", "WhisperDecoderLayer"]
 
     def _init_weights(self, module):
         std = self.config.init_std