Minor fixes for saving checkpoints. (#265)

* Minor fixes for saving checkpoints.

* Fix loading checkpoints saved by previous code.
This commit is contained in:
Fangjun Kuang 2022-03-23 12:22:05 +08:00 committed by GitHub
parent 8c7995d493
commit 6a091da0b0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -395,8 +395,9 @@ def load_checkpoint_if_available(
"cur_batch_idx",
]
for k in keys:
params[k] = saved_params[k]
params[k] = saved_params.get(k, 0)
if "cur_epoch" in saved_params:
params["start_epoch"] = saved_params["cur_epoch"]
return saved_params