Update Save Checkpoint

This commit is contained in:
cafeai 2022-12-03 20:06:46 +09:00
parent 34715bcc97
commit bf264d0ff0
1 changed files with 1 additions and 1 deletions

View File

@ -866,7 +866,7 @@ def main():
ema_unet.store(unet.parameters()) ema_unet.store(unet.parameters())
ema_unet.copy_to(unet.parameters()) ema_unet.copy_to(unet.parameters())
pipeline = StableDiffusionPipeline( pipeline = StableDiffusionPipeline(
text_encoder=text_encoder, text_encoder=text_encoder if type(text_encoder) is not torch.nn.parallel.DistributedDataParallel else text_encoder.module,
vae=vae, vae=vae,
unet=unet.module, unet=unet.module,
tokenizer=tokenizer, tokenizer=tokenizer,