Fixed errors

app.py

@@ -37,18 +37,22 @@ transform = transforms.Compose([
 
 # Prediction function
 def predict(image):
-    image = transform(image).unsqueeze(0)  # Add batch dimension
-    with torch.no_grad():
-        output = model(image)
-        probabilities = nn.Softmax(dim=1)(output)
-        predicted_class = torch.argmax(probabilities, dim=1)
-    return {str(i): probabilities[0][i].item() for i in range(10)}
+    try:
+        image = transform(image).unsqueeze(0)  # Add batch dimension
+        with torch.no_grad():
+            output = model(image)
+            probabilities = nn.Softmax(dim=1)(output)
+            predicted_class = torch.argmax(probabilities, dim=1)
+        return {str(i): probabilities[0][i].item() for i in range(10)}
+    except Exception as e:
+        print(f"Error in predict function: {e}")
+        return {"error": str(e)}
 
 # Create the Gradio interface
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Sketchpad(),
-    outputs=gr.Label(num_top_classes=10)
+    outputs=gr.Label()
 )
 
 # Launch the interface