Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 import torchvision.transforms as transforms
 from PIL import Image
-from transformers import AutoModel, AutoTokenizer
+from transformers import ViTForImageClassification, AutoTokenizer
 
 
 # Nombre del modelo en el repositorio de Hugging Face
@@ -10,7 +10,7 @@ model_name = "mkjaramillo/cancer"
 
 # Cargar el tokenizer y el modelo desde el repositorio de Hugging Face
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+model = ViTForImageClassification.from_pretrained(model_name)
 
 # Transformación de la imagen
 image_transform = transforms.Compose([