duzx16 commited on
Commit
047c657
1 Parent(s): 27f8067

Fix past_key_values

Browse files
Files changed (1) hide show
  1. modeling_chatglm.py +2 -1
modeling_chatglm.py CHANGED
@@ -932,9 +932,10 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
932
  standardize_cache_format: bool = False,
933
  ) -> Dict[str, Any]:
934
  # update past_key_values
935
- model_kwargs["past_key_values"] = self._extract_past_from_model_output(
936
  outputs, standardize_cache_format=standardize_cache_format
937
  )
 
938
 
939
  # update attention mask
940
  if "attention_mask" in model_kwargs:
 
932
  standardize_cache_format: bool = False,
933
  ) -> Dict[str, Any]:
934
  # update past_key_values
935
+ cache_name, cache = self._extract_past_from_model_output(
936
  outputs, standardize_cache_format=standardize_cache_format
937
  )
938
+ model_kwargs[cache_name] = cache
939
 
940
  # update attention mask
941
  if "attention_mask" in model_kwargs: