fixed def predict

app.py

@@ -46,21 +46,15 @@ transform = transforms.Compose([
 
 # --- Prediction function ---
 def predict(image):
-    """
-    Takes any image (JPG, PNG, etc.), converts to RGB, resizes to 32x32,
-    runs through the CNN, and returns class probabilities.
-    """
-    # Convert to RGB (in case of grayscale or RGBA input)
     image = image.convert("RGB")
-    image = transform(image).unsqueeze(0)  # shape: [1, 3, 32, 32]
-    
+    x = transform(image).unsqueeze(0)  # (1, 3, 32, 32)
     with torch.no_grad():
-        outputs=gr.Label(num_top_classes=3)
-        probs = torch.softmax(outputs, dim=1)[0]
-    
-    # Convert to dictionary: {class: probability}
+        outputs = model(x)              # tensor shape [1, 10]
+        probs = torch.nn.functional.softmax(outputs, dim=1)  # apply softmax
+        probs = probs[0].cpu().numpy()   # convert to numpy for Gradio
     return {classes[i]: float(probs[i]) for i in range(10)}
 
+
 # --- Gradio Interface ---
 demo = gr.Interface(
     fn=predict,