update modeling_baichuan.py for torchscript mode with past_kv

modeling_baichuan.py

@@ -285,7 +285,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:
@@ -297,6 +297,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
 
         seq_length_with_past = seq_length
@@ -437,7 +439,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]: