diff --git a/examples/tensorflow/simple.yaml b/examples/tensorflow/simple.yaml
index 4c2a0a76e6..f4c2fdd808 100644
--- a/examples/tensorflow/simple.yaml
+++ b/examples/tensorflow/simple.yaml
@@ -12,7 +12,9 @@ spec:
         spec:
           containers:
             - name: tensorflow
-              image: kubeflow/tf-mnist-with-summaries:latest
+# Source for the image is located here: https://github.com/kubeflow/katib/tree/master/examples/v1beta1/trial-images/tf-mnist-with-summaries
+              image: docker.io/kubeflowkatib/tf-mnist-with-summaries:latest
               command:
-                - "python"
-                - "/var/tf_mnist/mnist_with_summaries.py"
+              - "python"
+              - "/opt/tf-mnist-with-summaries/mnist.py"
+              - "--epochs=1"