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

Update trainer

pull/518/head
rainatam 2023-04-10 18:32:40 +08:00
parent 47a5ec121e
commit 2a5250ffcb
1 changed files with 5 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
ptuning/trainer.py
View File

@ -2825,12 +2825,11 @@ class Trainer:
state_dict = self.model.state_dict()
torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
else:
# state_dict = self.model.state_dict()
# filtered_state_dict = {}
# for k, v in self.model.named_parameters():
# if v.requires_grad:
# filtered_state_dict[k] = state_dict[k]
# print(filtered_state_dict.keys())
state_dict = self.model.state_dict()
filtered_state_dict = {}
for k, v in self.model.named_parameters():
if v.requires_grad:
filtered_state_dict[k] = state_dict[k]
self.model.save_pretrained(output_dir, state_dict=state_dict)
if self.tokenizer is not None:
self.tokenizer.save_pretrained(output_dir)