tensorflow_distributed_mnist_neo_inf1.ipynb failed CI

Link to the notebook:
https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker_neo_compilation_jobs/deploy_tensorflow_model_on_Inf1_instance/tensorflow_distributed_mnist_neo_inf1.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'

Start with sagemaker_neo_compilation_jobs/deploy_tensorflow_model_on_Inf1_instance/tensorflow_distributed_mnist_neo_inf1.ipynb and inspect the failing In [2] cell. Check the TensorFlow version used by CI and the replacement for the missing tensorflow.contrib MNIST import. Done means the notebook runs past that cell without ModuleNotFoundError and CI succeeds.

Written by the indexing model from the issue text.