Remove unsqueeze

HabanaAI · Sep 9, 2024 · 9916b6b · 9916b6b
1 parent 0440fb2
commit 9916b6b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vllm/worker/habana_model_runner.py b/vllm/worker/habana_model_runner.py
@@ -899,7 +899,7 @@ def _prepare_decode(
             for seq_id in seq_ids:
                 seq_data = seq_group_metadata.seq_data[seq_id]
                 generation_token = seq_data.get_last_token_id()
-                input_tokens.append(generation_token)
+                input_tokens.append([generation_token])
 
                 seq_len = seq_data.get_len()
                 position = seq_len - 1
@@ -928,7 +928,7 @@ def _prepare_decode(
             lora_logits_mask = lora_mask
         input_tokens = torch.tensor(input_tokens,
                                     dtype=torch.long,
-                                    device=self.device).unsqueeze(-1)
+                                    device=self.device)
         input_positions = torch.tensor(input_positions,
                                        dtype=torch.long,
                                        device=self.device)