keep logits.device unchanged

lvhan028 · Sep 24, 2024 · 28b0258 · 28b0258
1 parent e49b964
commit 28b0258
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lmdeploy/serve/utils.py b/lmdeploy/serve/utils.py
@@ -209,12 +209,12 @@ def get_ppl(self, inputs: List[str]) -> List[float]:
             ]
             steps = [i] * bs
             logits = generator.decode(
-                token_ids,
+                input_ids=token_ids,
                 steps=steps,
                 sequence_start=(i == 0),
                 sequence_end=(i + max_input_len >= max_seq_len))
             bsz, seq_len, vocab_size = logits.shape
-            logits = logits.float().cpu()
+            logits = logits.float()
             padding_token_id = -100
             # meaning logits[..., :, :] corresponds to labels
             # token_ids[1:] + predict_token_id, which is