2424

visualization.py

@@ -55,7 +55,8 @@ def vis(args,input_img=None,id=294,return_img=False):
 
     if args.frozen_weights is not None:
         print("Freeze weights for detector")
-
+    if not torch.cuda.is_available():
+        args.device = 'cpu'
     device = torch.device(args.device)
 
     # fix the seed for reproducibility