add --device for vae

vae/README.md

@@ -12,10 +12,11 @@ The main.py script accepts the following arguments:
 
 ```bash
 optional arguments:
-  --batch-size		input batch size for training (default: 128)
-  --epochs		number of epochs to train (default: 10)
-  --no-cuda		enables CUDA training
-  --mps         enables GPU on macOS
-  --seed		random seed (default: 1)
-  --log-interval	how many batches to wait before logging training status
-```
+  --batch-size N        input batch size for training (default: 128)
+  --epochs EPOCHS       number of epochs to train (default: 10)
+  --no-cuda             disables CUDA training
+  --no-mps              disables macOS GPU training
+  --device DEVICE       backend name
+  --seed SEED           random seed (default: 1)
+  --log-interval N      how many batches to wait before logging training status
+```

vae/main.py

@@ -17,6 +17,8 @@
                     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('--seed', type=int, default=1, metavar='S',
                     help='random seed (default: 1)')
 parser.add_argument('--log-interval', type=int, default=10, metavar='N',
@@ -32,7 +34,7 @@
 elif use_mps:
     device = torch.device("mps")
 else:
-    device = torch.device("cpu")
+    device = torch.device(args.device)
 
 kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}
 train_loader = torch.utils.data.DataLoader(