add --device for mnist_rnn

mnist_rnn/main.py

@@ -95,6 +95,8 @@ def main():
                         help='enables CUDA training')
     parser.add_argument('--mps', action="store_true", default=False,
                         help="enables MPS 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',
@@ -110,7 +112,7 @@ def main():
     elif args.mps and not args.cuda:
         device = "mps"
     else:
-        device = "cpu"
+        device = args.device
 
     device = torch.device(device)