From 88325c3bbcee6b913009022e269dd4089f85cae2 Mon Sep 17 00:00:00 2001
From: Laurent <laurent@lagon.tech>
Date: Fri, 12 Jul 2024 12:05:17 +0000
Subject: [PATCH] multi-upscaler: specify map_location when loading negative
 embedding

---
 .../latent_diffusion/stable_diffusion_1/multi_upscaler.py     | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/multi_upscaler.py b/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/multi_upscaler.py
index 4fb4c83..160c78d 100644
--- a/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/multi_upscaler.py
+++ b/src/refiners/foundationals/latent_diffusion/stable_diffusion_1/multi_upscaler.py
@@ -112,7 +112,9 @@ class MultiUpscalerAbstract(MultiDiffusion[T], ABC):
         if path is None:
             return ""
 
-        embeddings: Tensor | dict[str, Any] = torch.load(path, weights_only=True)  # type: ignore
+        embeddings: torch.Tensor | dict[str, Any] = torch.load(  # type: ignore
+            path, weights_only=True, map_location=self.device
+        )
 
         if isinstance(embeddings, dict):
             assert key is not None, "Key must be provided to access the negative embedding."