[wenet ]recognize.py fix dataset returns dict (#2111)

wenet-e2e · Nov 4, 2023 · 30c98dd · 30c98dd
1 parent 3342541
commit 30c98dd
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/wenet/bin/recognize.py b/wenet/bin/recognize.py
@@ -249,11 +249,11 @@ def main():
     max_format_len = max([len(mode) for mode in args.modes])
     with torch.no_grad():
         for batch_idx, batch in enumerate(test_data_loader):
-            keys, feats, target, feats_lengths, target_lengths = batch
-            feats = feats.to(device)
-            target = target.to(device)
-            feats_lengths = feats_lengths.to(device)
-            target_lengths = target_lengths.to(device)
+            keys = batch["keys"]
+            feats = batch["feats"].to(device)
+            target = batch["target"].to(device)
+            feats_lengths = batch["feats_lengths"].to(device)
+            target_lengths = batch["target_lengths"].to(device)
             results = model.decode(
                 args.modes,
                 feats,