From c37e04ea492d3a18862692f098a3e5e0577f4633 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=BF=97=E6=B5=A9?= <neo.dzh@alibaba-inc.com>
Date: Tue, 24 Sep 2024 17:45:55 +0800
Subject: [PATCH] add batch support for token extraction

---
 funasr/models/sense_voice/model_small.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/funasr/models/sense_voice/model_small.py b/funasr/models/sense_voice/model_small.py
index 41555cfad..0d27b4d25 100644
--- a/funasr/models/sense_voice/model_small.py
+++ b/funasr/models/sense_voice/model_small.py
@@ -2060,6 +2060,6 @@ class SenseVoiceL(nn.Module):
             ark_writer = self.writer
         if ark_writer is not None:
             for k, v, l in zip(key, tokens.detach().cpu().numpy(), out_lens):
-                ark_writer(k, tokens[:l])
+                ark_writer(k, v[:l])
 
         return results, meta_data