@@ -438,6 +438,7 @@ def load_model_hook(models, input_dir):
             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)):