diff --git a/pyproject.toml b/pyproject.toml
index aa994c4..b39794a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [project]
 name = "comfyui-unload-model"
 description = "For unloading a model or all models, using the memory management that is already present in ComfyUI. Copied from [a/https://github.com/willblaschko/ComfyUI-Unload-Models](https://github.com/willblaschko/ComfyUI-Unload-Models) but without the unnecessary extra stuff."
-version = "1.0.1"
+version = "1.0.2"
 license = {file = "LICENSE"}
 
 [project.urls]
diff --git a/unloadModel.py b/unloadModel.py
index d5f66e1..021c2cb 100644
--- a/unloadModel.py
+++ b/unloadModel.py
@@ -51,8 +51,7 @@ class UnloadModelNode:
             torch.cuda.ipc_collect()
         except:
             print("   - Unable to clear cache")
-        #time.sleep(2) # why?
-        return (list(kwargs.values()))
+        return (kwargs.get("value"),)
 
 class UnloadAllModelsNode:
     @classmethod
@@ -81,8 +80,7 @@ class UnloadAllModelsNode:
             torch.cuda.ipc_collect()
         except:
             print("   - Unable to clear cache")
-        #time.sleep(2) # why?
-        return (list(kwargs.values()))
+        return (kwargs.get("value"),)
 
 
 NODE_CLASS_MAPPINGS = {