1 files changed, 4 insertions, 4 deletions

diff --git a/dreambooth.py b/dreambooth.py
index 7b34fce..79b3d2c 100644
--- a/dreambooth.py
+++ b/dreambooth.py
@@ -489,7 +489,7 @@ class Checkpointer:
             generator=generator,
         )
 
-        with torch.inference_mode():
+        with torch.autocast("cuda"), torch.inference_mode():
             for pool, data, latents in [("stable", val_data, stable_latents), ("val", val_data, None), ("train", train_data, None)]:
                 all_samples = []
                 file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
@@ -522,7 +522,7 @@ class Checkpointer:
                         negative_prompt=nprompt,
                         height=self.sample_image_size,
                         width=self.sample_image_size,
-                        latents=latents[:len(prompt)] if latents is not None else None,
+                        latents_or_image=latents[:len(prompt)] if latents is not None else None,
                         generator=generator if latents is not None else None,
                         guidance_scale=guidance_scale,
                         eta=eta,
@@ -768,7 +768,7 @@ def main():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for batch in batched_data:
                     image_name = [item.class_image_path for item in batch]
                     prompt = [item.prompt.format(identifier=args.class_identifier) for item in batch]
@@ -1018,7 +1018,7 @@ def main():
             text_encoder.eval()
             val_loss = 0.0
 
-            with torch.inference_mode():
+            with torch.autocast("cuda"), torch.inference_mode():
                 for step, batch in enumerate(val_dataloader):
                     latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
                     latents = latents * 0.18215