Update app.py

app.py

@@ -11,12 +11,13 @@ Fruits = ['Acerola', 'Apple', 'Apricot', 'Avocado', 'Banana', 'Black Berry', 'Bl
           'Passion Fruit', 'Peach', 'Pear', 'Pineapple', 'Plum', 'Pomegranate', 'Raspberry', 'Strawberry', 'Tomato',
           'Watermelon']
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
 
 repo_name  = "VinayHajare/fruits30-resnet18"
 file_name = "fruit_resnet18(99.40%).pt"
 model_path = hf_hub_download(repo_id = repo_name, filename = file_name)
-model = torch.load(model_path).to(device)
+model = torch.load(model_path, map_location=torch.device('cpu'))
+model.to(device)
 model.eval()
 
 print("Stage 1: Completed ")
@@ -29,7 +30,8 @@ transform = transforms.Compose([
 ])
 
 def predict_image(image):
-    image_tensor = transform(Image.fromarray(image)).unsqueeze(0).to(device)
+    image_tensor = transform(Image.fromarray(image)).unsqueeze(0)
+    image_tensor = image_tensor.to(device)
     with torch.no_grad():
         output = model(image_tensor)
         predicted = torch.argmax(output).item()