Upload modeling_proteinglm.py

modeling_proteinglm.py

@@ -1276,7 +1276,8 @@ class ProteinGLMForCasualLM(ProteinGLMPreTrainedModel):
             is_encoder_decoder: bool = False,
     ) -> Dict[str, Any]:
         # update past_key_values
-        model_kwargs["past_key_values"] = self._extract_past_from_model_output(outputs)
+        cache_name, cache = self._extract_past_from_model_output(outputs)
+        model_kwargs[cache_name] = cache
 
         # update attention mask
         if "attention_mask" in model_kwargs: