Upload model_management.py

comfy/model_management.py

@@ -177,17 +177,14 @@ def get_torch_device():
     if cpu_state == CPUState.CPU:
         return torch.device("cpu")
     else:
-        if torch.cuda.is_available():  # Добавьте эту проверку!
-            if is_intel_xpu():
-                return torch.device("xpu", torch.xpu.current_device())
-            elif is_ascend_npu():
-                return torch.device("npu", torch.npu.current_device())
-            elif is_mlu():
-                return torch.device("mlu", torch.mlu.current_device())
-            else:
-                return torch.device(torch.cuda.current_device())
+        if is_intel_xpu():
+            return torch.device("xpu", torch.xpu.current_device())
+        elif is_ascend_npu():
+            return torch.device("npu", torch.npu.current_device())
+        elif is_mlu():
+            return torch.device("mlu", torch.mlu.current_device())
         else:
-            return torch.device("cpu")  # Fallback на CPU, если CUDA недоступен
+            return torch.device(torch.cuda.current_device())
 
 def get_total_memory(dev=None, torch_total_too=False):
     global directml_enabled