Spaces:

mmek
/

olive-classifier

Sleeping

mmek commited on Sep 10, 2024

Commit

a516a70

1 Parent(s): db88d52

change model load

Files changed (1) hide show

app.py CHANGED Viewed

@@ -4,7 +4,7 @@ import torch
 model = timm.create_model("mobileone_s2", pretrained = False)
 model.head.fc = torch.nn.Linear(model.head.fc.in_features,3)
-model.load_state_dict(torch.load("olive-classifier.pth", weights_only=True))
 model.eval()
 categories = ("Aculus Olearius", "Healthy", "Peacock Spot")

 model = timm.create_model("mobileone_s2", pretrained = False)
 model.head.fc = torch.nn.Linear(model.head.fc.in_features,3)
+model.load_state_dict(torch.load("olive-classifier.pth", map_location=torch.device('cpu'), weights_only=True))
 model.eval()
 categories = ("Aculus Olearius", "Healthy", "Peacock Spot")