Update app.py

app.py

@@ -11,14 +11,18 @@ def load_model():
     num_ftrs = model.classifier[6].in_features
     model.classifier[6] = nn.Linear(num_ftrs, len(class_names))  # Adjust this for your number of classes
 
-    # Load the model weights
-    model.load_state_dict(torch.load('model_alexnet.pth'))
-    model.eval()  # Set to evaluation mode
-
+    # Correctly map the model to CPU if CUDA is not available
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    if device == torch.device('cpu'):
+        model.load_state_dict(torch.load('model_alexnet.pth', map_location=device))
+    else:
+        model.load_state_dict(torch.load('model_alexnet.pth'))
+
+    model.eval()  # Set to evaluation mode
     model.to(device)
     return model, device
 
+
 model_alexnet, device = load_model()
 
 # Image transformations