Update app.py

app.py

@@ -51,13 +51,13 @@ transform = transforms.Compose([
     transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
 ])
 
-# -------- PREDICTION FUNCTION --------
-def predict_age(image):
-    image = transform(image).unsqueeze(0).to(device)
+def predict_age(image: Image.Image) -> float:
+    image_tensor = transform(image).unsqueeze(0).to(device)
     with torch.no_grad():
-        output = model(image)
-        age = output.item()  # Convert to a single scalar
-    return f"Predicted Age: {age:.2f}"
+        output = model(image_tensor)
+        age = output.item()
+    return round(age, 2)
+
 
 import gradio as gr