rm cuda

app.py

@@ -27,7 +27,7 @@ from model import SASNet
 warnings.filterwarnings('ignore')
 
 # define the GPU id to be used
-os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+#os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 
 class data(Dataset):
     def __init__(self, img, transform=None):
@@ -63,7 +63,8 @@ def loading_data(img):
 def predict(img):
     """the main process of inference"""
     test_loader = loading_data(img)
-    model = SASNet().cuda()
+    model = SASNet()
+    #model = SASNet().cuda()
     model_path = "./SHHA.pth"
     # load the trained model
     model.load_state_dict(torch.load(model_path))
@@ -74,9 +75,9 @@ def predict(img):
 
         for vi, data in enumerate(test_loader, 0):
           img = data
-          img = img.cuda()
+          #img = img.cuda()
           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