Update train_dreambooth_lora.py

train_dreambooth_lora.py

@@ -940,11 +940,11 @@ def main(args):
                 torch_dtype=weight_dtype,
             )
             pipeline.scheduler = DPMSolverMultistepScheduler.from_config(pipeline.scheduler.config)
-            pipeline = pipeline.to('cpu')
+            pipeline = pipeline.to(accelerator.device)
             pipeline.set_progress_bar_config(disable=True)
 
             # run inference
-            generator = torch.Generator(device='cpu').manual_seed(args.seed)
+            generator = torch.Generator(device=accelerator.device).manual_seed(args.seed)
             prompt = args.num_validation_images * [args.validation_prompt]
             images = pipeline(prompt, num_inference_steps=25, generator=generator).images