Skip to content

GitLab

Commit ac02b1b1 authored by Yizhou Wang's avatar Yizhou Wang
Browse files

Update train.py

parent 52716c59
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
tools/train.py
View file @ ac02b1b1
...@@ -280,8 +280,8 @@ if __name__ == "__main__": ...@@ -280,8 +280,8 @@ if __name__ == "__main__":
print("saving current model ...") print("saving current model ...")
status_dict = { status_dict = {
'model_name': model_name, 'model_name': model_name,
'epoch': epoch, 'epoch': epoch + 1,
'iter': iter, 'iter': iter + 1,
'model_state_dict': rodnet.state_dict(), 'model_state_dict': rodnet.state_dict(),
'optimizer_state_dict': optimizer.state_dict(), 'optimizer_state_dict': optimizer.state_dict(),
'loss': loss_confmap.item(), 'loss': loss_confmap.item(),
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment