fix

app.py

@@ -15,22 +15,23 @@ MODEL_LIST = {
 # Global variables
 current_model = None
 current_preprocessor = None
+device = "cuda" if torch.cuda.is_available() else "cpu"  # Dynamically set device
 
 # Load model and preprocessor
 def load_model_and_preprocessor(model_name):
     """Load model and preprocessor for a given model name."""
     global current_model, current_preprocessor
-    print(f"Loading model and preprocessor for: {model_name}")
-    current_model = AutoModelForImageClassification.from_pretrained(MODEL_LIST[model_name]).cuda().eval()
+    print(f"Loading model and preprocessor for: {model_name} on {device}")
+    current_model = AutoModelForImageClassification.from_pretrained(MODEL_LIST[model_name]).to(device).eval()
     current_preprocessor = AutoFeatureExtractor.from_pretrained(MODEL_LIST[model_name])
-    return f"Model {model_name} loaded successfully."
+    return f"Model {model_name} loaded successfully on {device}."
 
 # Predict function
 def predict(image, model, preprocessor):
     """Make a prediction on the given image patch using the loaded model."""
     if model is None or preprocessor is None:
         raise ValueError("Model and preprocessor are not loaded.")
-    inputs = preprocessor(images=image, return_tensors="pt").to("cuda")
+    inputs = preprocessor(images=image, return_tensors="pt").to(device)
     with torch.no_grad():
         outputs = model(**inputs)
         predicted_class = torch.argmax(outputs.logits, dim=1).item()