diff --git a/utils.py b/utils.py index 0b2b07e..fd2df67 100644 --- a/utils.py +++ b/utils.py @@ -29,7 +29,9 @@ def load_checkpoint(config, model, optimizer, lr_scheduler, logger): if not config.EVAL_MODE and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint: optimizer.load_state_dict(checkpoint['optimizer']) lr_scheduler.load_state_dict(checkpoint['lr_scheduler']) + config.defrost() config.TRAIN.START_EPOCH = checkpoint['epoch'] + 1 + config.freeze() if 'amp' in checkpoint and config.AMP_OPT_LEVEL != "O0" and checkpoint['config'].AMP_OPT_LEVEL != "O0": amp.load_state_dict(checkpoint['amp']) logger.info(f"=> loaded successfully '{config.MODEL.RESUME}' (epoch {checkpoint['epoch']})")