fix current length

wenet-e2e · Aug 14, 2024 · 5288f24 · 5288f24
1 parent e33da5f
commit 5288f24
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/wenet/ssl/init_dataset.py b/wenet/ssl/init_dataset.py
@@ -33,16 +33,21 @@ def __init__(
 
     def __iter__(self):
         for elem in self.dp:
-            current_length = self.length + self.length_fn(elem)
+            elem_length = self.length_fn(elem)
+            current_length = self.length + elem_length
             if current_length >= self.max_length:
                 long_elem = self.merge_fn(self.buf)
                 yield long_elem
                 del self.buf
                 self.buf = []
+                self.length = 0
             self.buf.append(elem)
+            self.length += current_length
         if len(self.buf) > 0:
             yield self.merge_fn(self.buf)
         del self.buf
+        self.buf = []
+        self.length = 0
 
 
 def cat_speech(buffer: List):