cpu

app.py

@@ -10,7 +10,7 @@ checkpoint_path = hf_hub_download(repo_id="ttoosi/resnet50_robust_face", filenam
 
 # Initialize the model
 model = models.resnet50()
-model.load_state_dict(torch.load(checkpoint_path))
+model.load_state_dict(torch.load(checkpoint_path, map_location=torch.device('cpu')))  # Force model to load on CPU
 model.eval()
 
 # Image preprocessing
@@ -25,7 +25,7 @@ preprocess = transforms.Compose([
 def predict(image):
     image = preprocess(image).unsqueeze(0)  # Add batch dimension
     with torch.no_grad():
-        output = model(image)
+        output = model(image)  # Perform inference on CPU
     _, predicted_class = output.max(1)
     return f"Predicted class: {predicted_class.item()}"