Reverse Huggingface model

source/predict_sample.py

@@ -3,7 +3,6 @@ import torch
 import torch.utils.data
 from PIL import Image
 from source.model import CNN
-from transformers import AutoModel
 
 
 def classify_eye(image: torch.Tensor,
@@ -43,9 +42,11 @@ def main_classification(image):
     image = transform(image)
     image = image.to(torch.device("cpu"))
 
-    cnn = AutoModel.from_pretrained("nssharmaofficial/RedEyeDetector")
+    cnn = CNN().to(torch.device("cpu"))
     cnn.eval()
 
+    cnn.load_state_dict(torch.load(f='source/weights/CNN-B8-LR-0.01-E30.pt', map_location=torch.device("cpu")))
+
     prediction_outcome = classify_eye(image, cnn)
 
     return prediction_outcome