use "comfy.model_management.get_torch_device()" instead if "torch.cuda.is_available"

2024-12-20 14:54:47 +02:00
parent 650cf53f08
commit 7fa2d0d357
1 changed files with 1 additions and 1 deletions
@@ -1453,7 +1453,7 @@ def create_mask_from_color_tensor(image:Image, color:str, tolerance:int=0) -> Im
 def load_RMBG_model():
    from .briarmbg import BriaRMBG
    current_directory = os.path.dirname(os.path.abspath(__file__))
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = comfy.model_management.get_torch_device()
    net = BriaRMBG()
    model_path = ""
    try: