Adjusting CPU for inferencing

app.py

@@ -4,9 +4,14 @@ from torchvision import models
 from PIL import Image
 import gradio as gr
 
+# Force CPU usage
+device = torch.device('cpu')
+
 # Load your trained ResNet-50 model
-model = models.resnet50(pretrained=False)  # Load the ResNet-50 architecture 
-model.load_state_dict(torch.load("model.pth"))  # Load the trained weights  (.pth)
+model = models.resnet50(pretrained=False)  # Load the ResNet-50 architecture
+model.load_state_dict(torch.load("model.pth", map_location=device))  # Load the trained weights (.pth)
+model.to(device)  # Move model to CPU (even if you have a GPU)
+
 model.eval()  # Set model to evaluation mode
 
 # Define the transformation required for the input image
@@ -18,7 +23,6 @@ transform = transforms.Compose([
 ])
 
 # Define the labels for ImageNet (or your specific dataset labels)
-# This is typically a list of class labels for classification
 LABELS = ["class_1", "class_2", "class_3", "class_4", "class_5",  # Replace with your classes
           "class_6", "class_7", "class_8", "class_9", "class_10"]
 
@@ -27,8 +31,12 @@ def predict(image):
     image = Image.open(image).convert("RGB")  # Open the image and convert to RGB
     image = transform(image).unsqueeze(0)  # Apply transformations and add batch dimension
 
+    # Move the image tensor to CPU as well
+    image = image.to(device)
+
     with torch.no_grad():
         outputs = model(image)  # Get model predictions
+
     _, predicted = torch.max(outputs, 1)  # Get the class with highest probability
     return LABELS[predicted.item()]  # Return the predicted class label