Former-commit-id: 0277d120e62164bb7fa1d6043b8fcc52c881fe96
This commit is contained in:
hiyouga 2023-11-05 12:21:07 +08:00
parent 6cfe1e1ac2
commit b2c3001f8e

View File

@ -196,6 +196,7 @@ def load_model_and_tokenizer(
# Fix LM head (for ChatGLM2 and ChatGLM3)
if getattr(config, "model_type", None) == "chatglm":
setattr(model, "lm_head", model.transformer.output_layer)
setattr(model, "_keys_to_ignore_on_save", ["lm_head.weight"])
# Register auto class to save the custom code files.
if isinstance(config, PretrainedConfig) and "AutoConfig" in getattr(config, "auto_map", {}):