fix lr resume for non-deepspeed ckpts

a2e7dabb · Zhang690683220 · 69a8a287 · a2e7dabb
Commit a2e7dabb authored Aug 02, 2022 by Zhang690683220
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

train_openfold.py train_openfold.py +5 -1

No files found.
--- a/train_openfold.py
+++ b/train_openfold.py
@@ -260,7 +260,11 @@ def main(args):
    
    model_module = OpenFoldWrapper(config)
    if(args.resume_from_ckpt):
-        last_global_step = get_global_step_from_zero_checkpoint(args.resume_from_ckpt)
+        if(os.path.isdir(args.resume_from_ckpt)):  
+            last_global_step = get_global_step_from_zero_checkpoint(args.resume_from_ckpt)
+        else:
+            sd = torch.load(args.resume_from_ckpt)
+            last_global_step = int(sd['global_step'])
        model_module.resume_last_lr_step(last_global_step)
        logging.info("Successfully loaded last lr step...")
    if(args.resume_from_ckpt and args.resume_model_weights_only):