2424

main.py

@@ -47,6 +47,8 @@ def main(args):
     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