fix mnist script

Author: Chuyang Deng

tests/data/pytorch_mnist/mnist.py

@@ -44,7 +44,7 @@ def _get_train_data_loader(training_dir, is_distributed, batch_size, **kwargs):
     dataset = datasets.MNIST(training_dir, train=True, transform=transforms.Compose([
         transforms.ToTensor(),
         transforms.Normalize((0.1307,), (0.3081,))
-    ]))
+    ]), download=True)
     train_sampler = torch.utils.data.distributed.DistributedSampler(dataset) if is_distributed else None
     train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=train_sampler is None,
                                                sampler=train_sampler, **kwargs)
@@ -57,7 +57,7 @@ def _get_test_data_loader(training_dir, **kwargs):
         datasets.MNIST(training_dir, train=False, transform=transforms.Compose([
             transforms.ToTensor(),
             transforms.Normalize((0.1307,), (0.3081,))
-        ])),
+        ]), download=True),
         batch_size=1000, shuffle=True, **kwargs)