github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[chat]fix save_model(#3377) 

The function save_model should be a part of PPOTrainer.
pull/3368/head
Dr-Corgi 2023-04-06 11:19:14 +08:00 committed by GitHub
parent 57a3c4db6d
commit 73afb63594
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
applications/Chat/coati/trainer/ppo.py
View File

@ -116,6 +116,9 @@ class PPOTrainer(Trainer):
self.critic_optim.zero_grad()
return {'reward': experience.reward.mean().item()}
def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)
def save_model(self, path: str, only_rank0: bool = False, tokenizer: Optional[PreTrainedTokenizerBase] = None) -> None:
self.strategy.save_model(model=self.actor, path=path, only_rank0=only_rank0, tokenizer=tokenizer)