pytorch-mnist-batch-transform.ipynb failed CI

Link to the notebook:
https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker_batch_transform/pytorch_mnist_batch_transform/pytorch-mnist-batch-transform.ipynb

Error:

---------------------------------------------------------------------------
Exception encountered at "In [2]":
---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
in
18 hyperparameters={
19 'epochs': 6,
---> 20 'backend': 'gloo'
21 })
22

/opt/conda/lib/python3.6/site-packages/sagemaker/pytorch/estimator.py in __init__(self, entry_point, framework_version, py_version, source_dir, hyperparameters, image_uri, distribution, **kwargs)
151 :class:`~sagemaker.estimator.EstimatorBase`.
152 """
--> 153 validate_version_or_image_args(framework_version, py_version, image_ur

[...]

ValueError: framework_version or py_version was None, yet image_uri was also None. Either specify both framework_version and py_version, or specify image_uri.

Open sagemaker_batch_transform/pytorch_mnist_batch_transform/pytorch-mnist-batch-transform.ipynb and inspect the setup in In [2], where the PyTorch estimator raises the reported ValueError. Run the notebook through the failing cell or the repository's CI validation, and confirm that it completes without this framework, Python-version, or image configuration error.

Written by the indexing model from the issue text.