multi-gpu inference? about bonito HOT 6 CLOSED

+++ b/bonito/train.py
@@ -69,6 +69,12 @@ def main(args):
             print("* error: Cannot use AMP: Apex package needs to be installed manually, See https://github.com/NVIDIA/apex")
             exit(1)
 
+    if args.multi-gpu:
+        from torch.nn import DataParallel
+        model = DataParallel(model)
+        model.stride = model.module.stride
+        model.alphabet = model.module.alphabet
+
     schedular = CosineAnnealingLR(optimizer, args.epochs * len(train_loader))
 
     log_interval = np.floor(len(train_dataset) / args.batch * 0.05)
@@ -80,7 +86,15 @@ def main(args):
             log_interval, model, device, train_loader, optimizer, epoch, use_amp=args.amp
         )
         test_loss, mean, median = test(model, device, test_loader)
+
+        if args.multi_gpu:
+            state = model.module.state_dict()
+        else:
+            state = model.state_dict()
+
+        # save optim state
         torch.save(model.state_dict(), os.path.join(workdir, "weights_%s.tar" % epoch))
+
         with open(os.path.join(workdir, 'training.csv'), 'a', newline='') as csvfile:
             csvw = csv.writer(csvfile, delimiter=',')
             if epoch == 1:
@@ -111,6 +125,7 @@ def argparser():
     parser.add_argument("--batch", default=32, type=int)
     parser.add_argument("--chunks", default=1000000, type=int)
     parser.add_argument("--validation_split", default=0.99, type=float)
+    parser.add_argument("--multi-gpu", action="store_true", default=False)
     parser.add_argument("--amp", action="store_true", default=False)
     parser.add_argument("-f", "--force", action="store_true", default=False)
     return parser