Spaces:

anon5
/

k-department-2025-02-01

Running

anon5 commited on 9 days ago

Commit

a8d4562

verified ·

1 Parent(s): 0576b0e

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -14,7 +14,7 @@ device = torch.device('cpu')
 model = models.resnet18(pretrained=True)
 model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)
-model.load_state_dict(torch.load('best_model.pth'))
 model.to(device)
 model.eval()

 model = models.resnet18(pretrained=True)
 model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)
+model.load_state_dict(torch.load('best_model.pth', map_location=device))
 model.to(device)
 model.eval()