From 2cf8b2a453963e3b4da69ed857514e058cca5a0a Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Sun, 20 Nov 2022 00:43:38 +0530
Subject: [PATCH] Use the correct device name when moving the model to cpu

---
 ui/sd_internal/runtime.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/ui/sd_internal/runtime.py b/ui/sd_internal/runtime.py
index d0efff6a..e217965d 100644
--- a/ui/sd_internal/runtime.py
+++ b/ui/sd_internal/runtime.py
@@ -245,9 +245,10 @@ def unload_models():
 
 def move_to_cpu(model):
     if thread_data.device != "cpu":
-        mem = torch.cuda.memory_allocated() / 1e6
+        d = torch.device(thread_data.device)
+        mem = torch.cuda.memory_allocated(d) / 1e6
         model.to("cpu")
-        while torch.cuda.memory_allocated() / 1e6 >= mem:
+        while torch.cuda.memory_allocated(d) / 1e6 >= mem:
             time.sleep(1)
 
 def load_model_gfpgan():