Hoist the autocast out of the loop.

Author: LucasSloan

examples/text_to_image/train_text_to_image.py

@@ -732,10 +732,10 @@ def collate_fn(examples):
                             # run inference
                             prompt = [args.validation_prompt]
                             images = []
-                            for _ in range(args.num_validation_images):
-                                with torch.autocast(
-                                    str(accelerator.device), enabled=accelerator.mixed_precision == "fp16"
-                                ):
+                            with torch.autocast(
+                                str(accelerator.device), enabled=accelerator.mixed_precision == "fp16"
+                            ):
+                                for _ in range(args.num_validation_images):
                                     images.append(pipeline(prompt).images[0])
 
                             for i, image in enumerate(images):