assert the model runs on gpu if available, if not, on cpu

app.py

@@ -15,7 +15,7 @@ from options.options import parse
 path_opt = './options/predict/LOLBlur.yml'
 
 opt = parse(path_opt)
-
+device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 #define some auxiliary functions
 pil_to_tensor = transforms.ToTensor()
 
@@ -29,12 +29,10 @@ model = Network(img_channel=opt['network']['img_channels'],
                     dilations=opt['network']['dilations'],
                     extra_depth_wise = opt['network']['extra_depth_wise'])
 
-checkpoints = torch.load(opt['save']['best'])
+checkpoints = torch.load(opt['save']['best'], map_location=device)
 # print(checkpoints)
 model.load_state_dict(checkpoints['model_state_dict'])
 
-
-device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 model = model.to(device)
 
 def load_img (filename):