Skip to content

Commit 2bfb097

Browse files
authoredJul 27, 2024··
chore: optimize tensor padding in model_runner.py (#639)
1 parent 319d037 commit 2bfb097

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed
 

‎ChatTTS/model/velocity/model_runner.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -569,13 +569,16 @@ def execute_model(
569569
for i in range(idx_next.shape[0]):
570570
idx_next_i = idx_next[i, 0, :].cpu().tolist()
571571
logprob_i = logprob[i].cpu().tolist()
572+
tmp_hidden_states = hidden_states[i].cpu()
573+
if input_tokens[i].shape[-2] != 1:
574+
tmp_hidden_states = tmp_hidden_states[-1:,:]
572575
result = SequenceGroupOutput(
573576
samples=[
574577
SequenceOutput(
575578
parent_seq_id=seq_groups[i],
576579
logprobs={tuple(idx_next_i): logprob_i},
577580
output_token=tuple(idx_next_i),
578-
hidden_states=hidden_states[i].cpu(),
581+
hidden_states=tmp_hidden_states,
579582
finished=finish[i].item(),
580583
),
581584
],

0 commit comments

Comments
 (0)
Please sign in to comment.