zRzRzRzRzRzRzR commited on
Commit
341d817
1 Parent(s): 24304a6

Update modeling_chatglm.py

Browse files
Files changed (1) hide show
  1. modeling_chatglm.py +2 -3
modeling_chatglm.py CHANGED
@@ -870,9 +870,8 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
870
  standardize_cache_format: bool = False,
871
  ) -> Dict[str, Any]:
872
  # update past_key_values
873
- model_kwargs["past_key_values"] = self._extract_past_from_model_output(
874
- outputs, standardize_cache_format=standardize_cache_format
875
- )
876
 
877
  # update attention mask
878
  if "attention_mask" in model_kwargs:
 
870
  standardize_cache_format: bool = False,
871
  ) -> Dict[str, Any]:
872
  # update past_key_values
873
+ cache_name, cache = self._extract_past_from_model_output(outputs)
874
+ model_kwargs[cache_name] = cache
 
875
 
876
  # update attention mask
877
  if "attention_mask" in model_kwargs: