Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
OpenDAS
ColossalAI
Commits
73afb635
Unverified
Commit
73afb635
authored
Apr 06, 2023
by
Dr-Corgi
Committed by
GitHub
Apr 06, 2023
Browse files
[chat]fix save_model(#3377)
The function save_model should be a part of PPOTrainer.
parent
57a3c4db
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
3 additions
and
0 deletions
+3
-0
applications/Chat/coati/trainer/ppo.py
applications/Chat/coati/trainer/ppo.py
+3
-0
No files found.
applications/Chat/coati/trainer/ppo.py
View file @
73afb635
...
@@ -116,6 +116,9 @@ class PPOTrainer(Trainer):
...
@@ -116,6 +116,9 @@ class PPOTrainer(Trainer):
self
.
critic_optim
.
zero_grad
()
self
.
critic_optim
.
zero_grad
()
return
{
'reward'
:
experience
.
reward
.
mean
().
item
()}
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
:
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
)
self
.
strategy
.
save_model
(
model
=
self
.
actor
,
path
=
path
,
only_rank0
=
only_rank0
,
tokenizer
=
tokenizer
)
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment