Skip to content

Commit c46669b

Browse files
authored
[BugFix] Remove redundant device_put. (#1099)
1 parent f817453 commit c46669b

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

tpu_inference/runner/tpu_runner.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1315,10 +1315,10 @@ def _prepare_inputs_dp(self, scheduler_output: "VllmSchedulerOutput"):
13151315
seq_lens_cpu = seq_lens
13161316

13171317
(input_ids, positions, block_tables, query_start_loc, seq_lens,
1318-
logits_indices, request_distribution, logits_indices) = device_array(
1318+
logits_indices, request_distribution) = device_array(
13191319
self.mesh,
13201320
(input_ids, positions, block_tables, query_start_loc, seq_lens,
1321-
logits_indices, request_distribution, logits_indices),
1321+
logits_indices, request_distribution),
13221322
sharding=data_parallel_attn_sharding,
13231323
)
13241324
# Async scheduling: substitute placeholder tokens for DP

0 commit comments

Comments
 (0)