fix doc of TextDecoder (#1526)

Signed-off-by: haoshengqiang <haoshengqiang@xiaohongshu.com> Co-authored-by: haoshengqiang <haoshengqiang@xiaohongshu.com>
openai · Sep 18, 2023 · 21010ef · 21010ef
1 parent 29b7df6
commit 21010ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/whisper/model.py b/whisper/model.py
@@ -197,7 +197,7 @@ def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = None):
         """
         x : torch.LongTensor, shape = (batch_size, <= n_ctx)
             the text tokens
-        xa : torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)
+        xa : torch.Tensor, shape = (batch_size, n_audio_ctx, n_audio_state)
             the encoded audio features to be attended on
         """
         offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0