From 602c63ce766199fbbc6fed4a2af267ef63280bf0 Mon Sep 17 00:00:00 2001
From: Laurent <laurent@lagon.tech>
Date: Sun, 8 Sep 2024 13:47:35 +0000
Subject: [PATCH] add missing device and dtype to SD1UNet's UpBlocks

---
 .../foundationals/latent_diffusion/stable_diffusion_1/unet.py   | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/unet.py b/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/unet.py
index 25e7502..54365ac 100644
--- a/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/unet.py
+++ b/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/unet.py
@@ -188,7 +188,7 @@ class SD1UNet(fl.Chain):
                 fl.UseContext(context="unet", key="residuals").compose(lambda x: x[-1]),
                 MiddleBlock(device=device, dtype=dtype),
             ),
-            UpBlocks(),
+            UpBlocks(device=device, dtype=dtype),
             fl.Chain(
                 fl.GroupNorm(channels=320, num_groups=32, device=device, dtype=dtype),
                 fl.SiLU(),