Update src/components/face_shape_detection.py

src/components/face_shape_detection.py

@@ -17,7 +17,7 @@ class FaceShapeDetection:
         self.classes = ['Heart', 'Oblong', 'Oval', 'Round', 'Square']
         self.model.load_state_dict(torch.load("/home/ishwor/Desktop/MannequinToModel/artifacts/best_model.pth"))
         self.model.eval()
-        self.model.to("cuda")
+        self.model.to("cpu")
 
         self.transform = T.Compose([
             T.ToPILImage(),
@@ -42,7 +42,7 @@ class FaceShapeDetectionTransformer:
     def __init__(self):
         self.t_model = pipeline("image-classification", model="metadome/face_shape_classification",
                                 torch_dtype=torch.float16,
-                                device="cuda")
+                                device="cpu")
 
     def detect_face_shape(self, image):
         result = self.t_model(image)