Commit 75dee47c authored by wat3rBro's avatar wat3rBro Committed by Francisco Massa

fix crash caused by saving model when staring from max iter (#215)

parent bcb32120
...@@ -101,10 +101,9 @@ def do_train( ...@@ -101,10 +101,9 @@ def do_train(
memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0, memory=torch.cuda.max_memory_allocated() / 1024.0 / 1024.0,
) )
) )
if iteration % checkpoint_period == 0: if iteration % checkpoint_period == 0 or iteration == max_iter:
checkpointer.save("model_{:07d}".format(iteration), **arguments) checkpointer.save("model_{:07d}".format(iteration), **arguments)
checkpointer.save("model_{:07d}".format(iteration), **arguments)
total_training_time = time.time() - start_training_time total_training_time = time.time() - start_training_time
total_time_str = str(datetime.timedelta(seconds=total_training_time)) total_time_str = str(datetime.timedelta(seconds=total_training_time))
logger.info( logger.info(
......
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