From 9c686cd8257e71ad51389acce4f176af6c861505 Mon Sep 17 00:00:00 2001 From: Yuekai Zhang Date: Mon, 15 Jan 2024 14:12:59 +0800 Subject: [PATCH] add repeat audio back --- whisper/decoding.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/whisper/decoding.py b/whisper/decoding.py index 49485d0..8d14dbf 100644 --- a/whisper/decoding.py +++ b/whisper/decoding.py @@ -730,7 +730,8 @@ class DecodingTask: ) ] - # repeat text tensors by the group size, for beam search or best-of-n sampling + # repeat the audio & text tensors by the group size, for beam search or best-of-n sampling + audio_features = audio_features.repeat_interleave(self.n_group, dim=0) tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device) # call the main sampling loop