minor fix

This commit is contained in:
JinZr 2023-09-15 11:02:20 +08:00
parent bd053ca90a
commit c3dc8f005b

View File

@ -76,7 +76,7 @@ def save_checkpoint(
if isinstance(model, DDP):
model = model.module
checkpoint = {
"model": model.state_dict(),
"optimizer": optimizer.state_dict() if optimizer is not None else None,