diff --git a/egs/librispeech/ASR/pruned_transducer_stateless4/train.py b/egs/librispeech/ASR/pruned_transducer_stateless4/train.py index 4e16770cb..2e2f97d98 100755 --- a/egs/librispeech/ASR/pruned_transducer_stateless4/train.py +++ b/egs/librispeech/ASR/pruned_transducer_stateless4/train.py @@ -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