AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

[Examples] Make sure EMA works with any device (#2382)

Browse Source 
* Fix EMA

* up

* update
This commit is contained in:
Patrick von Platen
2023-02-16 22:27:47 +02:00
committed by GitHub
parent a60f5555f5
commit ca980fd0d1
1 changed files with 1 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
examples/text_to_image/train_text_to_image.py
View File

@@ -438,6 +438,7 @@ def main():
if args.use_ema:
load_model = EMAModel.from_pretrained(os.path.join(input_dir, "unet_ema"), UNet2DConditionModel)
ema_unet.load_state_dict(load_model.state_dict())
ema_unet.to(accelerator.device)
del load_model
for i in range(len(models)):