github
/
ChatGLM2-6B
mirror of https://github.com/THUDM/ChatGLM2-6B
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge dd4bb36244 into db9d52d6ff

pull/334/merge
Mark 2024-05-12 18:32:46 +08:00 committed by GitHub
parent db9d52d6ff dd4bb36244
commit 9ce95e7caf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 2 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
ptuning/main.py
View File

@ -112,7 +112,7 @@ def main():
if model_args.ptuning_checkpoint is not None:
# Evaluation
# Loading extra state dict of prefix encoder
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True, device='cuda')
prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin"))
new_prefix_state_dict = {}
for k, v in prefix_state_dict.items():
@ -120,7 +120,7 @@ def main():
new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v
model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict)
else:
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True)
model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True, device='cuda')
if model_args.quantization_bit is not None:
print(f"Quantized to {model_args.quantization_bit} bit")