Update app.py

app.py

@@ -10,11 +10,11 @@ CLASSES = ['guro', 'pigs', 'proofs', 'protyk', 'safe', 'shit']
 NUM_CLASSES = len(CLASSES)
 
 #device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-device = 'cpu'
+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.load_state_dict(torch.load('best_model.pth', location=device))
 model.to(device)
 model.eval()