Some temp code for loading old state dict

This commit is contained in:
Daniel Povey 2022-05-28 00:29:47 +08:00
parent d89cb53a3b
commit 828defb019

View File

@ -875,13 +875,16 @@ def run(rank, world_size, args):
optimizer = Cain(model.parameters(), lr=params.initial_lr,
max_eff_lr=params.initial_lr)
scheduler = Eden(optimizer, params.lr_batches, params.lr_epochs)
if checkpoints and "optimizer" in checkpoints:
logging.info("Loading optimizer state dict")
optimizer.load_state_dict(checkpoints["optimizer"])
for p in optimizer.param_groups:
p['max_eff_lr'] = params.initial_lr # TEMP!!
if (
checkpoints
and "scheduler" in checkpoints