Update app.py

app.py

@@ -63,8 +63,8 @@ def loading_data(img):
 def predict(img):
     """the main process of inference"""
     test_loader = loading_data(img)
-    model = SASNet()
-    #model = SASNet().cuda()
+    #model = SASNet()
+    model = SASNet().cpu()
     model_path = "./SHHA.pth"
     # load the trained model
     model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
@@ -76,8 +76,9 @@ def predict(img):
         for vi, data in enumerate(test_loader, 0):
           img = data
           #img = img.cuda()
+          img = img.cpu()
           pred_map = model(img)
-          #pred_map = pred_map.data.cpu().numpy()
+          pred_map = pred_map.data.cpu().numpy()
           for i_img in range(pred_map.shape[0]):
             pred_cnt = np.sum(pred_map[i_img]) / 1000