fixing issues with cpu

test.py

@@ -526,10 +526,10 @@ def get_crossentropy_loss(gt,pre):
     return entropy_loss
 
 def get_alpha_loss(predict, alpha, trimap):
-    weighted = torch.zeros(trimap.shape).cuda()
+    weighted = torch.zeros(trimap.shape).to(device)
     weighted[trimap == 128] = 1.
     alpha_f = alpha / 255.
-    alpha_f = alpha_f.cuda()
+    alpha_f = alpha_f.to(device)
     diff = predict - alpha_f
     diff = diff * weighted
     alpha_loss = torch.sqrt(diff ** 2 + 1e-12)
@@ -537,9 +537,9 @@ def get_alpha_loss(predict, alpha, trimap):
     return alpha_loss_weighted
 
 def get_alpha_loss_whole_img(predict, alpha):
-    weighted = torch.ones(alpha.shape).cuda()
+    weighted = torch.ones(alpha.shape).to(device)
     alpha_f = alpha / 255.
-    alpha_f = alpha_f.cuda()
+    alpha_f = alpha_f.to(device)
     diff = predict - alpha_f
     alpha_loss = torch.sqrt(diff ** 2 + 1e-12)
     alpha_loss = alpha_loss.sum()/(weighted.sum())
@@ -555,7 +555,7 @@ def build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=False):
     kernel = np.sum(gaussian(grid), axis=2)
     kernel /= np.sum(kernel)
     kernel = np.tile(kernel, (n_channels, 1, 1))
-    kernel = torch.FloatTensor(kernel[:, None, :, :]).cuda()
+    kernel = torch.FloatTensor(kernel[:, None, :, :]).to(device)
     return Variable(kernel, requires_grad=False)
 
 def conv_gauss(img, kernel):
@@ -576,10 +576,10 @@ def laplacian_pyramid(img, kernel, max_levels=5):
     return pyr
 
 def get_laplacian_loss(predict, alpha, trimap):
-    weighted = torch.zeros(trimap.shape).cuda()
+    weighted = torch.zeros(trimap.shape).to(device)
     weighted[trimap == 128] = 1.
     alpha_f = alpha / 255.
-    alpha_f = alpha_f.cuda()
+    alpha_f = alpha_f.to(device)
     alpha_f = alpha_f.clone()*weighted
     predict = predict.clone()*weighted
     gauss_kernel = build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=True)
@@ -590,7 +590,7 @@ def get_laplacian_loss(predict, alpha, trimap):
 
 def get_laplacian_loss_whole_img(predict, alpha):
     alpha_f = alpha / 255.
-    alpha_f = alpha_f.cuda()
+    alpha_f = alpha_f.to(device)
     gauss_kernel = build_gauss_kernel(size=5, sigma=1.0, n_channels=1, cuda=True)
     pyr_alpha  = laplacian_pyramid(alpha_f, gauss_kernel, 5)
     pyr_predict = laplacian_pyramid(predict, gauss_kernel, 5)
@@ -598,7 +598,7 @@ def get_laplacian_loss_whole_img(predict, alpha):
     return laplacian_loss
 
 def get_composition_loss_whole_img(img, alpha, fg, bg, predict):
-    weighted = torch.ones(alpha.shape).cuda()
+    weighted = torch.ones(alpha.shape).to(device)
     predict_3 = torch.cat((predict, predict, predict), 1)
     comp = predict_3 * fg + (1. - predict_3) * bg
     comp_loss = torch.sqrt((comp - img) ** 2 + 1e-12)
@@ -781,7 +781,7 @@ def inference_img(model, img):
         img=cv2.copyMakeBorder(img, 8-h%8, 0, 8-w%8, 0, cv2.BORDER_REFLECT)
     # print(img.shape)
 
-    tensor_img = torch.from_numpy(img).permute(2, 0, 1).cuda()
+    tensor_img = torch.from_numpy(img).permute(2, 0, 1).to(device)
     input_t = tensor_img
     input_t = input_t/255.0
     normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
@@ -839,7 +839,7 @@ def test_am2k(model):
         alpha = alpha[:,:,0] if alpha.ndim>2 else alpha
 
         with torch.no_grad():
-            torch.cuda.empty_cache()
+#             torch.cuda.empty_cache()
             predict = inference_img( model, img)
 
         
@@ -926,7 +926,7 @@ def test_p3m10k(model,dataset_choice, max_image=-1):
         trimap = trimap[:,:,0] if trimap.ndim>2 else trimap
         alpha = alpha[:,:,0] if alpha.ndim>2 else alpha
         with torch.no_grad():
-            torch.cuda.empty_cache()
+#             torch.cuda.empty_cache()
             start = time.time()