add `accelerate` support for `Whisper` (#19697)

af556a09 · Younes Belkada · GitHub · fb0bd7b7 · af556a09
Unverified Commit af556a09 authored Oct 18, 2022 by Younes Belkada Committed by GitHub Oct 18, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

src/transformers/models/whisper/modeling_whisper.py src/transformers/models/whisper/modeling_whisper.py +1 -0

No files found.
--- 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