Update main_test_SRMNet.py

main_test_SRMNet.py

@@ -34,12 +34,14 @@ def main():
     if len(files) == 0:
         raise Exception(f"No files found at {inp_dir}")
 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
     # Load corresponding models architecture and weights
     model = SRMNet()
-    model.cuda()
-
-    load_checkpoint(model, args.weights)
+    model = model.to(device)
     model.eval()
+    load_checkpoint(model, args.weights)
+    
 
     mul = 16
     for file_ in files: