diff --git a/__init__.py b/__init__.py
index d9446b4..8af29fc 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,6 +1,7 @@
 import torch
 import comfy.utils
 from .Pytorch_Retinaface.pytorch_retinaface import Pytorch_RetinaFace
+from comfy.model_management import get_torch_device
 
 class AutoCropFaces:
     def __init__(self):
@@ -67,7 +68,7 @@ class AutoCropFaces:
 
     def auto_crop_faces_in_image (self, image, max_number_of_faces, scale_factor, shift_factor, aspect_ratio, method='lanczos'): 
         image_255 = image * 255
-        rf = Pytorch_RetinaFace(top_k=50, keep_top_k=max_number_of_faces)
+        rf = Pytorch_RetinaFace(top_k=50, keep_top_k=max_number_of_faces, device=get_torch_device())
         dets = rf.detect_faces(image_255)
         cropped_faces, bbox_info = rf.center_and_crop_rescale(image, dets, scale_factor=scale_factor, shift_factor=shift_factor, aspect_ratio=aspect_ratio)