diff --git a/python/llm/src/ipex_llm/transformers/models/baichuan.py b/python/llm/src/ipex_llm/transformers/models/baichuan.py
index a4be00ba..1f39c0bd 100644
--- a/python/llm/src/ipex_llm/transformers/models/baichuan.py
+++ b/python/llm/src/ipex_llm/transformers/models/baichuan.py
@@ -97,7 +97,7 @@ def baichuan_model_7b_forward(
     )
     use_cache = use_cache if use_cache is not None else self.config.use_cache
 
-    return_dict = return_dict if return_dict is not None else self.config.use_return_dict    
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
     if use_cache:
         inputs = input_ids if input_ids is not None else inputs_embeds
@@ -164,9 +164,6 @@ def baichuan_model_7b_forward(
 
     use_compresskv = isinstance(past_key_values, DynamicCompressCache)
 
-    # if not past_key_values and not use_compresskv:
-    #     past_key_values = [None for _ in range(self.num_layers)]
-
     for idx, decoder_layer in enumerate(self.layers):
         if output_hidden_states:
             all_hidden_states += (hidden_states,)