update model code

app.py

@@ -259,7 +259,7 @@ def predict():
         return jsonify({"message": "Not an mri image", "confidence": 0.95, "saved_path": image_save_path})
     a, b = model.check_file(temp_path)
     image = Image.open(temp_path).convert("RGB")
-    output = model(transform(image).unsqueeze(0).to(device))
+    output = model(transform(image).unsqueeze(0))
     stage = output.item()
     if stage <= 2.0:
         stage = "Mild"