diff --git a/whisper_timestamped/transcribe.py b/whisper_timestamped/transcribe.py
index e7127ee..255c994 100644
--- a/whisper_timestamped/transcribe.py
+++ b/whisper_timestamped/transcribe.py
@@ -1038,6 +1038,7 @@ def hook_output_logits(layer, ins, outs):
         all_hooks.append(model.decoder.ln.register_forward_hook(hook_output_logits))
 
     try:
+        model.alignment_heads = alignment_heads # Avoid exception "AttributeError: 'WhisperUntied' object has no attribute 'alignment_heads'. Did you mean: 'set_alignment_heads'?""
         transcription = model.transcribe(audio, **whisper_options)
     finally:
         for hook in all_hooks: