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

[pre-commit.ci] auto fixes from pre-commit.com hooks 

for more information, see https://pre-commit.ci
pull/6007/head
pre-commit-ci[bot] 2024-08-12 11:29:21 +00:00 committed by Tong Li
parent ba80449d62
commit e6245485fc
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
applications/ColossalChat/coati/trainer/rm.py
View File

@ -48,7 +48,7 @@ class RewardModelTrainer(SLTrainer):
model: Any,
booster: Booster,
optimizer: Optimizer,
plugin: Plugin,
plugin: Plugin,
lr_scheduler: _LRScheduler,
tokenizer: PreTrainedTokenizerBase,
loss_fn: Optional[Callable] = None,
@ -60,7 +60,9 @@ class RewardModelTrainer(SLTrainer):
save_dir: str = None,
coordinator: DistCoordinator = None,
) -> None:
super().__init__(booster, max_epochs=max_epochs, model=model, optimizer=optimizer, plugin=plugin, start_epoch=start_epoch)
super().__init__(
booster, max_epochs=max_epochs, model=model, optimizer=optimizer, plugin=plugin, start_epoch=start_epoch
)
self.actor_scheduler = lr_scheduler
self.tokenizer = tokenizer
self.loss_fn = loss_fn if loss_fn is not None else LogSigLoss(beta=beta)