Update app.py

app.py

@@ -3,9 +3,6 @@ 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():
@@ -25,7 +22,7 @@ def preprocess(image):
 # Define the predict function
 def predict(image):
     # Preprocess the image
-    input_tensor = transforms(image, return_tensors='pt')
+    input_tensor = preprocess(image)
     
     # Make a prediction
     with torch.no_grad():
@@ -44,10 +41,11 @@ def predict(image):
     
     return predictions
 
+
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(),
+    inputs=gr.Image(preprocess),
     outputs=gr.Label(num_top_classes=4),
     live=True
 )