Update app.py

app.py

@@ -3,6 +3,9 @@ import torch.nn as nn
 import torchvision.transforms as transforms
 from PIL import Image
 import gradio as gr
+from transformers import ViTFeatureExtractor
+
+transforms = ViTFeatureExtractor.from_pretrained('nateraw/vit-age-classifier')
 
 # Load your trained model
 with torch.no_grad():
@@ -22,7 +25,7 @@ def preprocess(image):
 # Define the predict function
 def predict(image):
     # Preprocess the image
-    input_tensor = preprocess(image)
+    input_tensor = transforms(image, return_tensors='pt')
     
     # Make a prediction
     with torch.no_grad():
@@ -30,6 +33,7 @@ def predict(image):
     
     # Perform post-processing if needed (e.g., softmax for probabilities)
     # Replace this with your actual post-processing logic
+    print(output.logits.argmax(1).item())
     probabilities = torch.softmax(output, dim=1).squeeze().tolist()
     
     # Map the class indices to class labels