Skip to content

Commit

Permalink
[cli] fix bugs with uncertain output results (#2127)
Browse files Browse the repository at this point in the history
Co-authored-by: user01 <user01@user01deMacBook-Pro.local>
  • Loading branch information
cdliang11 and user01 authored Nov 7, 2023
1 parent fd3803b commit 3c7f291
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions wenet/cli/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ def __init__(self, model_dir: str, gpu: int = -1):
model_path = os.path.join(model_dir, 'final.zip')
units_path = os.path.join(model_dir, 'units.txt')
self.model = torch.jit.load(model_path)
self.model.eval()
if gpu >= 0:
device = 'cuda:{}'.format(gpu)
else:
Expand All @@ -53,6 +54,7 @@ def compute_feats(self, audio_file: str) -> torch.Tensor:
feats = feats.unsqueeze(0)
return feats

@torch.no_grad()
def _decode(self,
audio_file: str,
tokens_info: bool = False,
Expand Down

0 comments on commit 3c7f291

Please sign in to comment.