Better cuda detection.

test.py

@@ -20,7 +20,12 @@ output_dir = sys.argv[2]
 device = torch.device('cuda' if is_cuda() else 'cpu')
 
 model = arch.RRDB_Net(1, 1, 64, 23, gc=32, upscale=4, norm_type=None, act_type='leakyrelu', mode='CNA', res_scale=1, upsample_mode='upconv')
-model.load_state_dict(torch.load(model_path, map_location=torch.device('cuda' if is_cuda() else 'cpu')), strict=True)
+
+if is_cuda():
+    model.load_state_dict(torch.load(model_path), strict=True)
+else:
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')), strict=True)
+    
 model.eval()
 
 for k, v in model.named_parameters():