提交 ced10f20 编写于 作者: W wat3rBro 提交者: Francisco Massa

always save checkpoint as model_final.pth after reaching max_iter (#218)

上级 ef8a93b1
......@@ -101,8 +101,10 @@ def do_train(
memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
)
)
if iteration % checkpoint_period == 0 or iteration == max_iter:
if iteration % checkpoint_period == 0:
checkpointer.save("model_{:07d}".format(iteration), **arguments)
if iteration == max_iter:
checkpointer.save("model_final", **arguments)
total_training_time = time.time() - start_training_time
total_time_str = str(datetime.timedelta(seconds=total_training_time))
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册