fix memory leak: ensure hooks are cleaned up on exception in find_alignment

This commit is contained in:
Viraj 2025-10-27 19:34:28 -05:00
parent c0d2f624c0
commit 1aae730936

View File

@ -193,13 +193,14 @@ def find_alignment(
from .model import disable_sdpa
try:
with torch.no_grad(), disable_sdpa():
logits = model(mel.unsqueeze(0), tokens.unsqueeze(0))[0]
sampled_logits = logits[len(tokenizer.sot_sequence) :, : tokenizer.eot]
token_probs = sampled_logits.softmax(dim=-1)
text_token_probs = token_probs[np.arange(len(text_tokens)), text_tokens]
text_token_probs = text_token_probs.tolist()
finally:
for hook in hooks:
hook.remove()