From 9b0ed8cd1ef524bc86e97b5e8b1cc0eceeceef2c Mon Sep 17 00:00:00 2001
From: josephrocca <1167575+josephrocca@users.noreply.github.com>
Date: Sat, 24 May 2025 07:11:44 +0800
Subject: [PATCH] fix: use torch device class instead of string for
 compatibility with comfy memory management

---
 nodes/models/flux.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/nodes/models/flux.py b/nodes/models/flux.py
index bc3b074..37767be 100644
--- a/nodes/models/flux.py
+++ b/nodes/models/flux.py
@@ -260,7 +260,7 @@ def load_model(
         data_type: str,
         **kwargs,
     ) -> tuple[FluxTransformer2DModel]:
-        device = f"cuda:{device_id}"
+        device = torch.device(f"cuda:{device_id}")
         prefixes = folder_paths.folder_names_and_paths["diffusion_models"][0]
         for prefix in prefixes:
             if os.path.exists(os.path.join(prefix, model_path)):