Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -46,7 +46,7 @@ def main():
     mul = 16
     for file_ in 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]