add accelerate support for Whisper (#19697)

2022-10-18 18:25:49 +02:00
parent fb0bd7b7a8
commit af556a09f6
1 changed files with 1 additions and 0 deletions
--- a/src/transformers/models/whisper/modeling_whisper.py
+++ b/src/transformers/models/whisper/modeling_whisper.py
@@ -446,6 +446,7 @@ class WhisperPreTrainedModel(PreTrainedModel):
    base_model_prefix = "model"
    main_input_name = "input_features"
    supports_gradient_checkpointing = True
+    _no_split_modules = ["WhisperEncoderLayer"]

    def _init_weights(self, module):
        std = self.config.init_std