From 4b66cf1126591cb6a7e0736c8e26aebbe154a78f Mon Sep 17 00:00:00 2001
From: layerdiffusion <19834515+lllyasviel@users.noreply.github.com>
Date: Wed, 14 Aug 2024 20:45:58 -0700
Subject: [PATCH] fix possible OOM again

---
 backend/patcher/base.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/backend/patcher/base.py b/backend/patcher/base.py
index dbeed0fa..c127fb3e 100644
--- a/backend/patcher/base.py
+++ b/backend/patcher/base.py
@@ -262,6 +262,8 @@ class ModelPatcher:
                     assert weight.module is not None, 'BNB bad weight without parent layer!'
                     bnb_layer = weight.module
                     if weight.bnb_quantized:
+                        weight_original_device = weight.device
+
                         if device_to is not None:
                             assert device_to.type == 'cuda', 'BNB Must use CUDA!'
                             weight = weight.to(device_to)
@@ -270,6 +272,9 @@ class ModelPatcher:
 
                         from backend.operations_bnb import functional_dequantize_4bit
                         weight = functional_dequantize_4bit(weight)
+
+                        if device_to is None:
+                            weight = weight.to(device=weight_original_device)
                     else:
                         weight = weight.data