fix: update pytorch mnist script

Author: Chuyang Deng

src/sagemaker/pytorch/defaults.py

@@ -12,7 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
-PYTORCH_VERSION = "0.4"
+PYTORCH_VERSION = '1.1.0'
 """Default PyTorch version for when the framework version is not specified.
 The latest PyTorch version is 1.1.0, but the default version is no longer updated so as to not break existing workflows.
 """

tests/data/pytorch_mnist/mnist.py

@@ -17,6 +17,7 @@
 logger.addHandler(logging.StreamHandler(sys.stdout))
 
 
+
 class Net(nn.Module):
     # Based on https://github.com/pytorch/examples/blob/master/mnist/main.py
     def __init__(self):
@@ -47,6 +48,7 @@ def _get_train_data_loader(training_dir, is_distributed, batch_size, **kwargs):
         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
@@ -70,6 +72,7 @@ def _get_test_data_loader(training_dir, **kwargs):
             transform=transforms.Compose(
                 [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
             ),
+            download=True
         ),
         batch_size=1000,
         shuffle=True,