Update app.py

app.py

@@ -3,8 +3,12 @@ import torch
 from PIL import Image
 from torchvision import transforms
 
+# Set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load the model
-model = torch.load("transfer_balanced_learning_model.pth")
+model = torch.load("transfer_balanced_learning_model.pth", map_location=torch.device('cpu'))
+model = model.to(device)
 model.eval()
 
 # Define the transforms for input images
@@ -17,14 +21,13 @@ transform = transforms.Compose([
 def predict(image):
     # Preprocess the image
     image = transform(image)
-    image = image.unsqueeze(0)  
+    image = image.unsqueeze(0).to(device)
     with torch.no_grad():
         output = model(image)
         _, predicted = torch.max(output, 1)
         class_names = ["COVID", "Lung_Opacity", "No_Tumor", "Normal", "Tumor", "Viral_Pneumonia"]
         return class_names[predicted.item()]
 
-
 demo = gr.Interface(
     fn=predict,
     inputs="image",
@@ -34,4 +37,4 @@ demo = gr.Interface(
 )
 
 # Launch the Gradio app
-demo.launch()
+demo.launch()