mrq/vall-e
5
9
Fork 0
Code Issues 8 Pull Requests Packages Projects Releases Wiki Activity

fixed mirostat issue

Browse Source
This commit is contained in:
mrq 2023-10-10 18:09:49 -05:00
parent 99e980d323
commit 47b3077415
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
vall_e/models/base.py
View File

@ -595,7 +595,7 @@ class Base(nn.Module):
if beam_width > 1: if beam_width > 1:
candidates = top_k_logits_list( logits, beam_width ) candidates = top_k_logits_list( logits, beam_width )
res = [ torch.tensor(token, device=logits[batch].device, dtype=torch.int16).unsqueeze(dim=-1) for batch, token in candidates ] res = [ torch.tensor(token, device=logits[batch].device, dtype=torch.int16).unsqueeze(dim=-1) for batch, token in candidates ]
scores = [ logits[batch].flatten()[token] for batch, token in candidates ] scores = [ logits[batch].flatten()[token].to(logits[batch].device) for batch, token in candidates ]
return res, scores return res, scores
# and sample # and sample