tensorflow_distributed_mnist_neo.ipynb failed CI

Link to the notebook:
https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker_neo_compilation_jobs/tensorflow_distributed_mnist/tensorflow_distributed_mnist_neo.ipynb

Error:

---------------------------------------------------------------------------
Exception encountered at "In [2]":
---------------------------------------------------------------------------
ModuleNotFoundError Traceback (most recent call last)
in
1 import utils
----> 2 from tensorflow.contrib.learn.python.learn.datasets import mnist
3 import tensorflow as tf
4
5 data_sets = mnist.read_data_sets('data', dtype=tf.uint8, reshape=False, validation_size=5000)

ModuleNotFoundError: No module named 'tensorflow.contrib'

Open sagemaker_neo_compilation_jobs/tensorflow_distributed_mnist/tensorflow_distributed_mnist_neo.ipynb and inspect cell In[2], starting with the failing tensorflow.contrib.learn.datasets.mnist import. Run the notebook or its CI check and verify that the MNIST data-loading step completes without the reported ModuleNotFoundError.

Written by the indexing model from the issue text.