Skip to content

GitLab

Commit 98e1fe31 authored by Yoach Lacombe's avatar Yoach Lacombe
Browse files

fix ddp issue

parent e7cc576a
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 4 deletions
run_stable_speech_training.py
View file @ 98e1fe31
......@@ -1389,10 +1389,8 @@ def main():
if cur_step == total_train_steps:
# un-wrap student model for save
model = accelerator.unwrap_model(model)
model.save_pretrained(training_args.output_dir)
# re-wrap student model for final eval
model = accelerator.prepare(model)
unwrapped_model = accelerator.unwrap_model(model)
unwrapped_model.save_pretrained(training_args.output_dir)
if training_args.push_to_hub:
repo.push_to_hub(
......
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