Updated app.py and requirements.txt

app.py

@@ -3,11 +3,11 @@ import torch
 from PIL import Image
 from transformers import AutoModelForImageClassification, AutoImageProcessor
 
-# Load model and processor
-model = AutoModelForImageClassification.from_pretrained("shravvvv/SAG-ViT")
-processor = AutoImageProcessor.from_pretrained("shravvvv/SAG-ViT")
+# Load model and processor with custom code enabled
+model = AutoModelForImageClassification.from_pretrained("shravvvv/SAG-ViT", trust_remote_code=True)
+processor = AutoImageProcessor.from_pretrained("shravvvv/SAG-ViT", trust_remote_code=True)
 
-# CIFAR-10 class labels
+# Define CIFAR-10 class labels
 class_labels = [
     'airplane', 'automobile', 'bird', 'cat', 'deer',
     'dog', 'frog', 'horse', 'ship', 'truck'

requirements.txt

@@ -3,8 +3,6 @@ pandas==2.2.3
 matplotlib==3.7.5
 seaborn==0.12.2
 tqdm==4.66.4
-psutil==5.9.3
-pynvml==11.4.1
 scikit-learn==1.2.2
 torch==2.4.0
 torch-geometric==2.6.1