Small fix to prevent issues when using a gpu. (#228)

* Small fix to prevent issues when using a gpu.

* Relax the pytorch requirement to allow the 2.3.0 release

---------

Co-authored-by: Jose Sotelo <jose@jdescript.com>
Co-authored-by: Vicki Anand <vicki_anand@hotmail.com>

Author: 3 people

models/utils/detect_face.py

@@ -373,6 +373,11 @@ def extract_face(img, box, image_size=160, margin=0, save_path=None):
         os.makedirs(os.path.dirname(save_path) + "/", exist_ok=True)
         save_img(face, save_path)
 
-    face = F.to_tensor(np.float32(face))
+    if isinstance(face, np.ndarray) or isinstance(face, Image.Image):
+        face = F.to_tensor(np.float32(face))
+    elif isinstance(face, torch.Tensor):
+        face = face.float()
+    else:
+        raise NotImplementedError
 
     return face

setup.py

@@ -39,8 +39,8 @@
         'numpy>=1.24.0,<2.0.0',
         'Pillow>=10.2.0,<10.3.0',
         'requests>=2.0.0,<3.0.0',
-        'torch>=2.2.0,<2.3.0',
-        'torchvision>=0.17.0,<0.18.0',
+        'torch>=2.2.0,<=2.3.0',
+        'torchvision>=0.17.0,<=0.18.0',
         'tqdm>=4.0.0,<5.0.0',
     ],
 )