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

* Fix EMA * up * update

[Examples] Make sure EMA works with any device (#2382)
* Fix EMA * up * update
ca980fd0 · Patrick von Platen · GitHub · a60f5555 · ca980fd0
Unverified Commit ca980fd0 authored Feb 16, 2023 by Patrick von Platen Committed by GitHub Feb 16, 2023
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

examples/text_to_image/train_text_to_image.py examples/text_to_image/train_text_to_image.py +1 -0

No files found.
--- a/examples/text_to_image/train_text_to_image.py
+++ b/examples/text_to_image/train_text_to_image.py
@@ -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)):