diff --git a/lmdeploy/pytorch/models/baichuan.py b/lmdeploy/pytorch/models/baichuan.py
index c775114b80..fd5ad68ab1 100644
--- a/lmdeploy/pytorch/models/baichuan.py
+++ b/lmdeploy/pytorch/models/baichuan.py
@@ -217,7 +217,7 @@ def _continuous_batching_forward_7b(
         output_hidden_states: Optional[bool] = None,
         return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
-        """Rewrite implementation of LlamaModel.forward."""
+        """Rewrite implementation of 7b BaichuanModel.forward."""
         output_attentions = (output_attentions if output_attentions is not None
                              else self.config.output_attentions)
         output_hidden_states = (output_hidden_states