diff --git a/nodes.py b/nodes.py
index e129988..061e2b9 100644
--- a/nodes.py
+++ b/nodes.py
@@ -176,7 +176,9 @@ class GGUFModelPatcher(comfy.model_patcher.ModelPatcher):
                 patches = getattr(p, "patches", [])
                 if len(patches) > 0:
                     p.patches = []
-        return super().unpatch_model(device_to=None, unpatch_weights=unpatch_weights)
+        # TODO: Find another way to not unload after patches
+        device_to = device_to if comfy.model_management.DISABLE_SMART_MEMORY else None
+        return super().unpatch_model(device_to=device_to, unpatch_weights=unpatch_weights)
 
     def load(self, *args, force_patch_weights=False, **kwargs):
         # always call `patch_weight_to_device` even for lowvram