def METHOD_NAME(num_workers=4):
    """ Partitioning MNIST """
    dataset = datasets.MNIST(root=local_dataset_path, train=True, download=True,
                             transform=transforms.Compose([
                                 transforms.ToTensor(),
                                 transforms.Normalize((0.1307,), (0.3081,))
                             ]))
    size = dist.get_world_size()
    bsz = int(128 / float(size))
    partition_sizes = [1.0 / size for _ in range(size)]
    partition = DataPartitioner(dataset, partition_sizes)
    partition = partition.use(dist.get_rank())
    train_set = th.utils.data.DataLoader(
        partition, num_workers=num_workers, batch_size=bsz, shuffle=True)
    return train_set, bsz