From 140374b557eaccb2e12a415b80accf7b4655a9a8 Mon Sep 17 00:00:00 2001 From: DKnight54 <126916963+DKnight54@users.noreply.github.com> Date: Sun, 30 Mar 2025 04:41:36 +0800 Subject: [PATCH] Update train_util.py --- library/train_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/library/train_util.py b/library/train_util.py index 1054986d..6aa0631e 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -5719,7 +5719,7 @@ def sample_image_inference( clean_memory_on_device(accelerator.device) - original_latent = prompt_dict.get("original_lantent") + original_latent = prompt_dict.get("original_lantent").to(device=accelerator.device) logger.info(f"original_latent: {original_latent.shape}") original_image = pipeline.latents_to_image(original_latent)[0] text_image = draw_text_on_image(f"caption: {prompt}", image.width * 2)