mirror of
https://github.com/openai/whisper.git
synced 2025-11-23 22:15:58 +00:00
fix memory leak: ensure hooks are cleaned up on exception in find_alignment
This commit is contained in:
parent
c0d2f624c0
commit
1aae730936
@ -193,13 +193,14 @@ def find_alignment(
|
|||||||
|
|
||||||
from .model import disable_sdpa
|
from .model import disable_sdpa
|
||||||
|
|
||||||
|
try:
|
||||||
with torch.no_grad(), disable_sdpa():
|
with torch.no_grad(), disable_sdpa():
|
||||||
logits = model(mel.unsqueeze(0), tokens.unsqueeze(0))[0]
|
logits = model(mel.unsqueeze(0), tokens.unsqueeze(0))[0]
|
||||||
sampled_logits = logits[len(tokenizer.sot_sequence) :, : tokenizer.eot]
|
sampled_logits = logits[len(tokenizer.sot_sequence) :, : tokenizer.eot]
|
||||||
token_probs = sampled_logits.softmax(dim=-1)
|
token_probs = sampled_logits.softmax(dim=-1)
|
||||||
text_token_probs = token_probs[np.arange(len(text_tokens)), text_tokens]
|
text_token_probs = token_probs[np.arange(len(text_tokens)), text_tokens]
|
||||||
text_token_probs = text_token_probs.tolist()
|
text_token_probs = text_token_probs.tolist()
|
||||||
|
finally:
|
||||||
for hook in hooks:
|
for hook in hooks:
|
||||||
hook.remove()
|
hook.remove()
|
||||||
|
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user