[Minor] Improve the style and fix flaky tests (#1584)

2024-10-06 00:10:48 -07:00
parent 2422de5193
commit 9244f27f0a
7 changed files with 15 additions and 7 deletions
--- a/python/sglang/srt/sampling/sampling_batch_info.py
+++ b/python/sglang/srt/sampling/sampling_batch_info.py
@@ -40,7 +40,7 @@ class SamplingBatchInfo:
    @classmethod
    def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
        reqs = batch.reqs
-        with torch.device("cuda"):
+        with batch.input_ids.device:
            temperatures = torch.tensor(
                [r.sampling_params.temperature for r in reqs],
                dtype=torch.float,