initial model push

README.md

@@ -1,13 +1,3 @@
----
-library_name: torchvision          # hangi framework’ü kullandığı
-pipeline_tag: image-classification # görev türü (widget için şart)
-tags:
-  - image-classification
-  - efficientnet
-  - garbage
-metrics:
-  - accuracy
----
 # Garbage Classifier · EfficientNet‑V2‑S (torchvision)
 
 Finetuned model for 10‑class garbage image classification.

handler.py

@@ -1,5 +1,6 @@
 from typing import Dict, Any
 import io, base64, torch, torchvision
+from safetensors.torch import load_file
 from PIL import Image
 from torchvision import transforms as T
 
@@ -14,7 +15,8 @@ class EndpointHandler:
             torch.nn.Dropout(0.5),
             torch.nn.Linear(256, len(self.labels))
         )
-        self.model.load_state_dict(torch.load(f"{path}/model.safetensors", map_location="cpu"))
+        state = load_file(str(pth), device="cpu")
+        self.model.load_state_dict(state)
         self.model.eval()
         self.trans = T.Compose([
             T.Resize((224, 224)),