//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
context = [mx.cpu()] if args.gpus is None or args.gpus == "" else \
[mx.gpu(int(x)) for x in args.gpus.split(",")]
assert args.batch_size % len(context) == 0, \
"Total batch size must be multiple of the number of devices"
After Change
// Load data
//////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
context = [mx.cpu()] if not args.gpu else [mx.gpu(int(args.gpu))]
assert args.batch_size % len(context) == 0, \
"Total batch size must be multiple of the number of devices"
