Add the files

requirements.txt

@@ -1,2 +1,4 @@
-transformers
-gradio
+transformers==4.29.0
+datasets==2.10.0
+torch==2.1.0
+gradio==3.33.0

train.py

@@ -11,7 +11,7 @@ feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16
 
 # Preprocessamento
 def preprocess_function(examples):
-    return feature_extractor(examples["image"], return_tensors="pt")
+    return feature_extractor(examples["img"], return_tensors="pt")
 
 # Aplicando o preprocessamento ao dataset
 dataset = dataset.map(preprocess_function, batched=True)