Update app.py

app.py

@@ -21,7 +21,7 @@ num_ftrs = model.classifier[-1].in_features  # Get the number of input features
 model.classifier[-1] = nn.Linear(num_ftrs, 2)  # Replace with a new linear layer for binary classification
 
 # Load the saved weights into the model
-model.load_state_dict(torch.load('rice_plant_classification.pth', weights_only=True))  # Ensure this file exists
+model.load_state_dict(torch.load('rice_plant_classification.pth', map_location=torch.device('cpu')))  # Ensure this file exists
 model.eval()
 
 # Define class labels