Unverified Commit 3c8694ea authored by Ximingwang-09's avatar Ximingwang-09 Committed by GitHub
Browse files

Fix some typo (#19475)


Signed-off-by: default avatarximing.wxm <ximing.wxm@antgroup.com>
Co-authored-by: default avatarximing.wxm <ximing.wxm@antgroup.com>
parent 7484e1fc
......@@ -283,14 +283,14 @@ class RejectionSampler(SpecDecodeStochasticBaseSampler):
batch_size, k, _ = draft_probs.shape
batch_indices = torch.arange(batch_size,
device=target_probs.device)[:, None]
probs_indicies = torch.arange(k, device=target_probs.device)
probs_indices = torch.arange(k, device=target_probs.device)
# shape [batch_size, k]
selected_draft_probs = draft_probs[batch_indices, probs_indicies,
selected_draft_probs = draft_probs[batch_indices, probs_indices,
draft_token_ids]
# shape [batch_size, k]
selected_target_probs = target_probs[batch_indices, probs_indicies,
selected_target_probs = target_probs[batch_indices, probs_indices,
draft_token_ids]
uniform_rand = self._create_uniform_samples(seeded_seqs, batch_size,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment