From fe470ae5ade3d4b04ec5ee6f26ffb681d66a8ca1 Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Date: Sat, 9 Dec 2023 19:24:29 -0800
Subject: [PATCH] [Minor] Fix code style for baichuan (#2003)

---
 vllm/model_executor/models/baichuan.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/vllm/model_executor/models/baichuan.py b/vllm/model_executor/models/baichuan.py
index f7a3b90a..4a07c395 100644
--- a/vllm/model_executor/models/baichuan.py
+++ b/vllm/model_executor/models/baichuan.py
@@ -366,8 +366,8 @@ class BaiChuanBaseForCausalLM(nn.Module):
                 weight_loader(param, loaded_weight)
 
 
-class BaichuanForCausalLM(BaiChuanBaseForCausalLM
-                          ):  # baichuan 13b, baichuan2 13b, baichuan2 7b
+class BaichuanForCausalLM(BaiChuanBaseForCausalLM):
+    """Baichuan 13B and Baichuan2 7B/13B."""
 
     def __init__(self,
                  config,
@@ -378,7 +378,8 @@ class BaichuanForCausalLM(BaiChuanBaseForCausalLM
             super().__init__(config, "ALIBI", linear_method)
 
 
-class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):  # baichuan 7b
+class BaiChuanForCausalLM(BaiChuanBaseForCausalLM):
+    """Baichuan 7B."""
 
     def __init__(self,
                  config,