fix

main.py

@@ -102,7 +102,7 @@ trainer_args.add_argument('--epochs', '-e', type=int, default=200,
     help='the number of epochs')
 trainer_args.add_argument('--obs-epochs', '-oe', type=int, default=5,
     help='the number of sub epochs for observation stage')
-trainer_args.add_argument('--batch-size', '-bs', type=int, default=1,
+trainer_args.add_argument('--batch-size', '-bs', type=int, default=32,
     help='input batch size for training (default: 1)')
 trainer_args.add_argument('--eval-batch-size', '-ebs', type=int, default=1,
     help='input batch size for evaluation (default: 32)')
@@ -450,8 +450,8 @@ def main():
     model = get_model(args)
     
     if args.use_gpu:
-        # model = DataParallel(model).cuda()
-        model = model.cuda()
+        model = DataParallel(model).cuda()
+        # model = model.cuda()
         
     if args.weight_decay == 0:
         optimizer = optim.Adam(model.parameters(), lr=args.lr)