Make input_ids a torch.Tensor (#1568)

2024-10-04 01:09:59 -07:00
parent 114bbc8651
commit 45473d4b2b
3 changed files with 11 additions and 7 deletions
--- a/python/sglang/srt/mem_cache/memory_pool.py
+++ b/python/sglang/srt/mem_cache/memory_pool.py
@@ -30,6 +30,7 @@ class ReqToTokenPool:
    def __init__(self, size: int, max_context_len: int, device: str):
        self.size = size
        self.max_context_len = max_context_len
+        self.device = device
        self.free_slots = list(range(size))
        self.req_to_token = torch.empty(
            (size, max_context_len), dtype=torch.int32, device=device