Skip to content
GitLab
Menu
Projects
Groups
Snippets
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Andrei-Claudiu Roibu
BrainMapper
Commits
a897de36
Commit
a897de36
authored
Jul 31, 2020
by
Andrei Roibu
Browse files
fixed checkpoint loader
parent
91ecfd99
Changes
1
Hide whitespace changes
Inline
Side-by-side
solver.py
View file @
a897de36
...
...
@@ -376,14 +376,14 @@ class Solver():
self
.
start_epoch
=
checkpoint
[
'epoch'
]
self
.
start_iteration
=
checkpoint
[
'start_iteration'
]
# We are not loading the model_name as we might want to pre-train a model and then use it.
self
.
model
.
load_state_dict
=
checkpoint
[
'state_dict'
]
self
.
optimizer
.
load_state_dict
=
checkpoint
[
'optimizer'
]
self
.
model
.
load_state_dict
(
checkpoint
[
'state_dict'
]
)
self
.
optimizer
.
load_state_dict
(
checkpoint
[
'optimizer'
]
)
for
state
in
self
.
optimizer
.
state
.
values
():
for
state
in
self
.
optimizer
.
state
.
():
for
key
,
value
in
state
.
items
():
if
torch
.
is_tensor
(
value
):
state
[
key
]
=
value
.
to
(
self
.
device
)
self
.
learning_rate_scheduler
.
load_state_dict
=
checkpoint
[
'scheduler'
]
self
.
learning_rate_scheduler
.
load_state_dict
(
checkpoint
[
'scheduler'
]
)
self
.
LogWriter
.
log
(
"Checkpoint Loaded {} - epoch {}"
.
format
(
checkpoint_file_path
,
checkpoint
[
'epoch'
]))
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