Update app.py

app.py

@@ -52,22 +52,18 @@ def predict(image: Image.Image) -> Dict[str, float]:
     """
     try:
         image = image.convert("RGB")
-        image_tensor = preprocess(image).unsqueeze(0)
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        image_tensor = image_tensor.to(device)
-        model.to(device)
+        image_tensor = preprocess(image).unsqueeze(0).to(model.device)  # Используем device модели
         
         with torch.no_grad():
             output = model(image_tensor)
             scores = torch.nn.functional.softmax(output[0], dim=0)
         
         predictions = {label: float(score) for label, score in zip(labels, scores)}
-        sorted_predictions = dict(sorted(predictions.items(), key=lambda item: item[1], reverse=True))
-        
-        return sorted_predictions
+        return dict(sorted(predictions.items(), key=lambda item: item[1], reverse=True))
     except Exception as e:
         return {"error": str(e)}
 
+
 title = "🔥 Skin-AI"
 description = (
     "🔬 **Skin-AI — AI-Powered Skin Disease Classification**\n\n"