Update app.py

app.py

@@ -7,7 +7,7 @@ from PIL import Image
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 model = torch.hub.load("pytorch/vision", "resnet101", pretrained=False)
 model.fc = nn.Sequential(nn.Linear(2048, 500), nn.ReLU(), nn.Linear(500, 2), nn.Softmax(1))
-state_dict = torch.load('model.pth', map_location=torch.device('cpu'))
+state_dict = torch.load('model.pth', map_location=device)
 model.load_state_dict(state_dict)
 model.to(device)
 model.eval()