fixes deepspeed function definition.

ea142a0a · Jennifer · 1ae833bf · ea142a0a
Commit ea142a0a authored Apr 19, 2024 by Jennifer
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 1 deletion

train_openfold.py train_openfold.py +1 -1

No files found.
--- a/train_openfold.py
+++ b/train_openfold.py
@@ -306,7 +306,7 @@ def main(args):
        if args.resume_model_weights_only:
            # Load the checkpoint
            if os.path.isdir(args.resume_from_ckpt):
-                sd = get_fp32_state_dict_from_zero_checkpoint(
+                sd = zero_to_fp32.get_fp32_state_dict_from_zero_checkpoint(
                    args.resume_from_ckpt)
            else:
                sd = torch.load(args.resume_from_ckpt)