def get_mnist_iterator(rank):
    data_dir = ('data-%d' % rank)
    if (not os.path.isdir(data_dir)):
        os.makedirs(data_dir)
    zip_file_path = download('http://data.mxnet.io/mxnet/data/mnist.zip', dirname=data_dir)
    with zipfile.ZipFile(zip_file_path) as zf:
        zf.extractall(data_dir)
    input_shape = (1, 28, 28)
    batch_size = args.batch_size
    train_iter = mx.io.MNISTIter(image=('%s/train-images-idx3-ubyte' % data_dir), label=('%s/train-labels-idx1-ubyte' % data_dir), input_shape=input_shape, batch_size=batch_size, shuffle=True, flat=False, num_parts=hvd.size(), part_index=hvd.rank())
    val_iter = mx.io.MNISTIter(image=('%s/t10k-images-idx3-ubyte' % data_dir), label=('%s/t10k-labels-idx1-ubyte' % data_dir), input_shape=input_shape, batch_size=batch_size, flat=False)
    return (train_iter, val_iter)
def get_mnist_iterator(rank):
    data_dir = ('data-%d' % rank)
    if (not os.path.isdir(data_dir)):
        os.makedirs(data_dir)
    zip_file_path = download('http://data.mxnet.io/mxnet/data/mnist.zip', dirname=data_dir)
    with zipfile.ZipFile(zip_file_path) as zf:
        zf.extractall(data_dir)
    input_shape = (1, 28, 28)
    batch_size = args.batch_size
    train_iter = mx.io.MNISTIter(image=('%s/train-images-idx3-ubyte' % data_dir), label=('%s/train-labels-idx1-ubyte' % data_dir), input_shape=input_shape, batch_size=batch_size, shuffle=True, flat=False, num_parts=hvd.size(), part_index=hvd.rank())
    val_iter = mx.io.MNISTIter(image=('%s/t10k-images-idx3-ubyte' % data_dir), label=('%s/t10k-labels-idx1-ubyte' % data_dir), input_shape=input_shape, batch_size=batch_size, flat=False)
    return (train_iter, val_iter)