Skip to content

GitLab

Commit 2cd5f18a authored by Mostofa Patwary's avatar Mostofa Patwary
Browse files

added reload_model_params function

parent 8eff2a99
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
megatron/training.py
View file @ 2cd5f18a
...@@ -304,7 +304,7 @@ def setup_model_and_optimizer(model_provider_func): ...@@ -304,7 +304,7 @@ def setup_model_and_optimizer(model_provider_func):
print_rank_0("Initializing ICT from pretrained BERT model") print_rank_0("Initializing ICT from pretrained BERT model")
unwrapped_model.init_state_dict_from_bert() unwrapped_model.init_state_dict_from_bert()
if args.fp16: if args.fp16:
optimizer._copy_model_params_to_main_params() optimizer.reload_model_params()
return model, optimizer, lr_scheduler return model, optimizer, lr_scheduler
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment