fix: use torch device class instead of string for compatibility with comfy memory management (#200)

Author: josephrocca

nodes/models/flux.py

@@ -272,7 +272,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)):