diff --git a/py/imagefunc.py b/py/imagefunc.py
index cb276d2..9c74e2a 100644
--- a/py/imagefunc.py
+++ b/py/imagefunc.py
@@ -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: