misc: correct the int data type for token ids and indices (#969)

2024-08-07 13:40:07 -07:00
parent 4d929107ae
commit 6db27f7b3b
2 changed files with 3 additions and 3 deletions
--- a/python/sglang/srt/managers/schedule_batch.py
+++ b/python/sglang/srt/managers/schedule_batch.py
@@ -780,7 +780,7 @@ def top_k_top_p_sampling_from_probs_torch(
        sampled_index = torch.multinomial(probs_sort, num_samples=1)
    except RuntimeError:
        batch_next_token_ids = torch.zeros(
-            (probs_sort.shape[0],), dtype=torch.int64, device=probs.device
+            (probs_sort.shape[0],), dtype=torch.int32, device=probs.device
        )
        success = torch.zeros(probs.shape[0], dtype=torch.bool, device=probs.device)
        return batch_next_token_ids, success