[Whisper] Fix gradient checkpointing (#19538)

bdfcbe60 · Sanchit Gandhi · GitHub · 4edb3e49 · bdfcbe60
Unverified Commit bdfcbe60 authored Oct 12, 2022 by Sanchit Gandhi Committed by GitHub Oct 12, 2022
Show 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
@@ -20,6 +20,7 @@ import random
 from typing import Optional, Tuple

 import torch
+import torch.utils.checkpoint
 from torch import nn
 from torch.nn import CrossEntropyLoss