add --device for siamese_network

pytorch · shink · Sep 20, 2024 · Sep 20, 2024 · Sep 20, 2024 · Sep 20, 2024
commit 61929977a6e1f3e3bd65b348fcbb3fb5699fb752
diff --git a/siamese_network/main.py b/siamese_network/main.py
@@ -251,6 +251,8 @@ def main():
                         help='disables CUDA training')
     parser.add_argument('--no-mps', action='store_true', default=False,
                         help='disables macOS GPU training')
+    parser.add_argument('--device', type=str, default='cpu',
+                        help='backend device')
     parser.add_argument('--dry-run', action='store_true', default=False,
                         help='quickly check a single pass')
     parser.add_argument('--seed', type=int, default=1, metavar='S',
@@ -271,7 +273,7 @@ def main():
     elif use_mps:
         device = torch.device("mps")
     else:
-        device = torch.device("cpu")
+        device = torch.device(args.device)
 
     train_kwargs = {'batch_size': args.batch_size}
     test_kwargs = {'batch_size': args.test_batch_size}