Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -45,7 +45,7 @@ def main():
 
     for i, file_ in enumerate(tqdm(files)):
         img = Image.open(file_).convert('RGB')
-        input_ = TF.to_tensor(img).unsqueeze(0).cuda()
+        input_ = TF.to_tensor(img).unsqueeze(0).to(device)
 
         # Pad the input if not_multiple_of 8
         h, w = input_.shape[2], input_.shape[3]