update modeling_baichuan.py for torchscript mode with past_kv

modeling_baichuan.py

@@ -365,7 +365,7 @@ class BaichuanModel(BaichuanPreTrainedModel):
         use_cache: Optional[bool] = False,
         output_attentions: Optional[bool] = False,
         output_hidden_states: Optional[bool] = False,
-        return_dict: Optional[bool] = True,
+        return_dict: Optional[bool] = None,
     ) -> Union[Tuple, BaseModelOutputWithPast]:
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError(
@@ -378,6 +378,8 @@ class BaichuanModel(BaichuanPreTrainedModel):
         else:
             raise ValueError("You need to provide input_ids or inputs_embeds")
 
+        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
         )
@@ -682,7 +684,7 @@ class BaichuanForCausalLM(BaichuanPreTrainedModel):
         use_cache: Optional[bool] = None,
         output_attentions: Optional[bool] = False,
         output_hidden_states: Optional[bool] = False,
-        return_dict: Optional[bool] = True,
+        return_dict: Optional[bool] = None,
         **kwargs,
     ) -> Union[Tuple, CausalLMOutputWithPast]:
         return_dict = (