fix inference mode error

irexyc · Oct 28, 2024 · 7334cd2 · 7334cd2
1 parent a8b2765
commit 7334cd2
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/lmdeploy/pytorch/models/qwen2_vl.py b/lmdeploy/pytorch/models/qwen2_vl.py
@@ -388,10 +388,11 @@ def forward(
             inputs_embeds=inputs_embeds,
             mrope_position_ids=mrope_position_ids,
         )
+        return hidden_states
 
-        logits = self.lm_head(hidden_states)
-        logits = logits.float()
-        return logits
+    def get_logits(self, hidden_states: torch.Tensor):
+        """compute logits of the model output."""
+        return self.lm_head(hidden_states)
 
     def update_weights(self):
         """update weights."""