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() state_dict = self.model.state_dict()
torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME))
else: else:
# state_dict = self.model.state_dict() state_dict = self.model.state_dict()
# filtered_state_dict = {} filtered_state_dict = {}
# for k, v in self.model.named_parameters(): for k, v in self.model.named_parameters():
# if v.requires_grad: if v.requires_grad:
# filtered_state_dict[k] = state_dict[k] filtered_state_dict[k] = state_dict[k]
# print(filtered_state_dict.keys())
self.model.save_pretrained(output_dir, state_dict=state_dict) self.model.save_pretrained(output_dir, state_dict=state_dict)
if self.tokenizer is not None: if self.tokenizer is not None:
self.tokenizer.save_pretrained(output_dir) self.tokenizer.save_pretrained(output_dir)