4w4kt

/

waste-classifier-b5

like 1

Fixed standalone inference snippet:

• The previous code rebuilt the classifier as a single Linear layer, but the trained model's classifier is Linear(2048,512) → ReLU → Dropout(0.5) → Linear(512, num_classes). The mismatch meant that load_state_dict would throw on load. Now the snippet reconstructs the real architecture.
• Switched loading from torch.load (model.pth) to safetensors.load_file (model.safetensors).

P.S.: .pth can execute code on load, so .safetensorsis safer for a public deployment. Could you please convert and upload the model as model.safetensors? Glad to share the conversion script if useful.