riaz.somc/transformers
1
0
Fork 0
mirror of https://github.com/huggingface/transformers.git synced 2025-08-03 03:31:05 +06:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

Add gradient_checkpointing parameter to FlaxWhisperEncoder (#23300)

Browse Source 
Add gradient_checkpointing parameter
This commit is contained in:
raghavanone 2023-05-11 23:43:05 +05:30 committed by GitHub
parent 83eda6435e
commit ab96bf0294
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/transformers/models/whisper/modeling_flax_whisper.py
View File

@ -1515,7 +1515,9 @@ class FlaxWhisperForAudioClassificationModule(nn.Module):
gradient_checkpointing: bool = False
def setup(self) -> None:
self.encoder = FlaxWhisperEncoder(config=self.config, dtype=self.dtype)
self.encoder = FlaxWhisperEncoder(
config=self.config, dtype=self.dtype, gradient_checkpointing=self.gradient_checkpointing
)
self.config.is_encoder_decoder = False
num_layers = self.config.num_hidden_layers + 1
if self.config.use_weighted_layer_sum: