HF checkpoints i removed the optimizer to save space, init Adam without the first/second moments is ok

2023-07-27 22:20:07 +00:00
parent cc66a2037e
commit 517763346d
1 changed files with 1 additions and 1 deletions
@@ -179,7 +179,7 @@ scaler = torch.cuda.amp.GradScaler(enabled=(dtype == "float16"))

 # optimizer
 optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
-if init_from == "resume":
+if init_from == "resume" and "optimizer" in checkpoint:
    optimizer.load_state_dict(checkpoint["optimizer"])
 checkpoint = None  # free up memory