From b2ea4424a3b285b51723a4cf3ed7548d5173d474 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B8=B8=E9=9B=81?= <zhifu.gzf@alibaba-inc.com>
Date: Wed, 21 Feb 2024 17:20:15 +0800
Subject: [PATCH] llm

---
 funasr/models/paraformer/model.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py
index 729b8f500..90ce162d4 100644
--- a/funasr/models/paraformer/model.py
+++ b/funasr/models/paraformer/model.py
@@ -455,7 +455,9 @@ class Paraformer(torch.nn.Module):
             speech, speech_lengths = data_in, data_lengths
             if len(speech.shape) < 3:
                 speech = speech[None, :, :]
-            if speech_lengths is None:
+            if speech_lengths is not None:
+                speech_lengths = speech_lengths.squeeze(-1)
+            else:
                 speech_lengths = speech.shape[1]
         else:
             # extract fbank feats