From 72c9ce79b1646f1b919e8e91653394b691256677 Mon Sep 17 00:00:00 2001 From: jianl Date: Sat, 7 Sep 2024 20:19:38 +0800 Subject: [PATCH] Fix:median_out not implemented for half --- lib_layerdiffusion/models.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lib_layerdiffusion/models.py b/lib_layerdiffusion/models.py index 19396f7..f192acd 100644 --- a/lib_layerdiffusion/models.py +++ b/lib_layerdiffusion/models.py @@ -305,7 +305,8 @@ def estimate_augmented(self, pixel, latent): in gpu vram with dimensions higher than 4, we move it to cpu, call torch.median() and then move the result back to gpu. ''' - median = torch.median(result.cpu(), dim=0).values + result_cpu = result.cpu().float() # Convert to float32 on CPU + median = torch.median(result_cpu, dim=0).values median = median.to(device=self.load_device, dtype=self.dtype) else: median = torch.median(result, dim=0).values