Fixed minor bug

when unloading vae during TI training, generating images after training will error out
3 years ago · 890e68aaf7
parent 3b0127e698
commit 890e68aaf7
1 changed files with 1 additions and 0 deletions
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@ -409,6 +409,7 @@ Last saved image: {html.escape(last_saved_image)}<br/>

    filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')
    save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True)
+    shared.sd_model.first_stage_model.to(devices.device)

    return embedding, filename