From ef55b8fc0fa0267b38b2abbdb0b5d7f2dc170646 Mon Sep 17 00:00:00 2001 From: Kinda Hall <71250227+Alndaly@users.noreply.github.com> Date: Fri, 5 Jul 2024 20:52:34 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20Fix=20the=20`torch.median?= =?UTF-8?q?()`=20on=20apple=20gpu=20(#94)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ✅ Closes: #1 --- lib_layerdiffusion/models.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/lib_layerdiffusion/models.py b/lib_layerdiffusion/models.py index c9e6d3c..19396f7 100644 --- a/lib_layerdiffusion/models.py +++ b/lib_layerdiffusion/models.py @@ -299,7 +299,16 @@ class TransparentVAEDecoder: result += [eps] result = torch.stack(result, dim=0) - median = torch.median(result, dim=0).values + if self.load_device == torch.device("mps"): + ''' + In case that apple silicon devices would crash when calling torch.median() on tensors + 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 + median = median.to(device=self.load_device, dtype=self.dtype) + else: + median = torch.median(result, dim=0).values return median @torch.no_grad()