From 65604ada8d40a764ae99af46092ac7e7756fedea Mon Sep 17 00:00:00 2001
From: u010280923 <xinlu.nlp@foxmail.com>
Date: Mon, 13 Mar 2023 15:44:48 +0800
Subject: [PATCH] opt reward model
---
src/trainer.py | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/src/trainer.py b/src/trainer.py
index b43becd..257d2a5 100644
--- a/src/trainer.py
+++ b/src/trainer.py
@@ -248,7 +248,7 @@ class rm_train_callback(pl.Callback):
to_save_dict = pl_module.state_dict()
my_save(
to_save_dict,
- f"{args.proj_dir}/rwkv-final.pth",
+ f"{args.proj_dir}/rm-final.pth",
)
@@ -276,7 +276,7 @@ class rm_train_callback(pl.Callback):
try:
my_save(
to_save_dict,
- f"{args.proj_dir}/rwkv-{args.epoch_begin + trainer.current_epoch}.pth",
+ f"{args.proj_dir}/rm-{args.epoch_begin + trainer.current_epoch}.pth",
)
except Exception as e:
print('Error\n\n', e, '\n\n')
--
GitLab