Fix KeyError when looking for 'softmax_tensor' in the layers tutorial.

PiperOrigin-RevId: 161247019

1 files changed, 8 insertions, 7 deletions

diff --git a/tensorflow/examples/tutorials/layers/cnn_mnist.py b/tensorflow/examples/tutorials/layers/cnn_mnist.py
index f92277dac7..2124843fcb 100644
--- a/tensorflow/examples/tutorials/layers/cnn_mnist.py
+++ b/tensorflow/examples/tutorials/layers/cnn_mnist.py
@@ -86,13 +86,14 @@ def cnn_model_fn(features, labels, mode):
   # Output Tensor Shape: [batch_size, 10]
   logits = tf.layers.dense(inputs=dropout, units=10)
 
-  # Generate Predictions (for PREDICT mode)
-  predicted_classes = tf.argmax(input=logits, axis=1)
+  predictions = {
+      # Generate predictions (for PREDICT and EVAL mode)
+      "classes": tf.argmax(input=logits, axis=1),
+      # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
+      # `logging_hook`.
+      "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
+  }
   if mode == tf.estimator.ModeKeys.PREDICT:
-    predictions = {
-        "classes": predicted_classes,
-        "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
-    }
     return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
 
   # Calculate Loss (for both TRAIN and EVAL modes)
@@ -111,7 +112,7 @@ def cnn_model_fn(features, labels, mode):
   # Add evaluation metrics (for EVAL mode)
   eval_metric_ops = {
       "accuracy": tf.metrics.accuracy(
-          labels=labels, predictions=predicted_classes)}
+          labels=labels, predictions=predictions["classes"])}
   return tf.estimator.EstimatorSpec(
       mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)