def load_mnist(is_train = True, flatten = True):
    from torchvision import datasets, transforms

    dataset = datasets.MNIST(
        '../data', train = is_train, download = True,
        transform = transforms.Compose([
            transforms.ToTensor(),
        ]),
    )

    x = dataset.data.float() / 255.
    y = dataset.targets

    if flatten:
        x = x.view(x.size(0), -1)

        return x, y