diff --git a/how-to-use-azureml/ml-frameworks/tensorflow/train-tensorflow-resume-training/train-tensorflow-resume-training.ipynb b/how-to-use-azureml/ml-frameworks/tensorflow/train-tensorflow-resume-training/train-tensorflow-resume-training.ipynb
index c3bcd342..430eccef 100644
--- a/how-to-use-azureml/ml-frameworks/tensorflow/train-tensorflow-resume-training/train-tensorflow-resume-training.ipynb
+++ b/how-to-use-azureml/ml-frameworks/tensorflow/train-tensorflow-resume-training/train-tensorflow-resume-training.ipynb
@@ -420,7 +420,9 @@
         "                      script='tf_mnist_with_checkpoint.py',\n",
         "                      arguments=args,\n",
         "                      compute_target=compute_target,\n",
-        "                      environment=tf_env)"
+        "                      environment=tf_env)",
+        "\n",
+        "src.run_config.data_references = {checkpoint_data_ref.data_reference_name : checkpoint_data_ref.to_config()}"
       ]
     },
     {
@@ -497,4 +499,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 2
-}
\ No newline at end of file
+}