Skip to content

GitLab

"git@developer.sourcefind.cn:chenpangpang/diffusers.git" did not exist on "b09a2aa308f9e0a9ec06cb82f1f2a7bccf21a787"
Unverified Commit ac84c2fa authored by Suraj Patil's avatar Suraj Patil Committed by GitHub
Browse files

[textual-inversion] fix saving embeds (#387) 

fix saving embeds
parent 5a38033d
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 1 deletion
examples/textual_inversion/textual_inversion.py
View file @ ac84c2fa
...@@ -564,7 +564,7 @@ def main(): ...@@ -564,7 +564,7 @@ def main():
pipeline.save_pretrained(args.output_dir) pipeline.save_pretrained(args.output_dir)
# Also save the newly trained embeddings # Also save the newly trained embeddings
learned_embeds = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[placeholder_token_id] learned_embeds = accelerator.unwrap_model(text_encoder).get_input_embeddings().weight[placeholder_token_id]
learned_embeds_dict = {args.placeholder_token: learned_embeds} learned_embeds_dict = {args.placeholder_token: learned_embeds.detach().cpu()}
torch.save(learned_embeds_dict, os.path.join(args.output_dir, "learned_embeds.bin")) torch.save(learned_embeds_dict, os.path.join(args.output_dir, "learned_embeds.bin"))
if args.push_to_hub: if args.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