Updated code to use safetensors

config.json

@@ -12,5 +12,6 @@
   "patch_size": [
     4,
     4
-  ]
+  ],
+  "use_safetensors": true
 }

test_model_loading.py

@@ -1,6 +1,11 @@
 from transformers import AutoModelForImageClassification, AutoImageProcessor
+from safetensors import safe_open
 
-model = AutoModelForImageClassification.from_pretrained("shravvvv/SAG-ViT", trust_remote_code=True)
-processor = AutoImageProcessor.from_pretrained("shravvvv/SAG-ViT", trust_remote_code=True)
+model = AutoModelForImageClassification.from_pretrained(
+    "shravvvv/SAG-ViT",
+    use_safetensors=True,  # Enable safetensors
+    trust_remote_code=True
+)
+processor = AutoImageProcessor.from_pretrained("shravvvv/SAG-ViT", use_safetensors=True, trust_remote_code=True)
 
 print(model)