"web/git@developer.sourcefind.cn:chenpangpang/ComfyUI.git" did not exist on "d35267e85a865c30a5fa63fdb0a21f94f4cc37e7"
Commit 856a63da authored by wangfei's avatar wangfei
Browse files

Fix: save model/model.module

parent 1ef41b83
...@@ -322,7 +322,8 @@ def main(): ...@@ -322,7 +322,8 @@ def main():
# Save a trained model # Save a trained model
if args.local_rank == -1 or torch.distributed.get_rank() == 0: if args.local_rank == -1 or torch.distributed.get_rank() == 0:
logging.info("** ** * Saving fine-tuned model ** ** * ") logging.info("** ** * Saving fine-tuned model ** ** * ")
model.save_pretrained(args.output_dir) model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir) tokenizer.save_pretrained(args.output_dir)
......
...@@ -610,7 +610,8 @@ def main(): ...@@ -610,7 +610,8 @@ def main():
# Save a trained model # Save a trained model
if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0): if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
logger.info("** ** * Saving fine - tuned model ** ** * ") logger.info("** ** * Saving fine - tuned model ** ** * ")
model.save_pretrained(args.output_dir) model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
model_to_save.save_pretrained(args.output_dir)
tokenizer.save_pretrained(args.output_dir) tokenizer.save_pretrained(args.output_dir)
......
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