Как получить вероятность каждого класса в примере переобучения TensorFlow?

Question

Этот является примером переобучения TensorFlow, и я хотел бы получить вероятность каждого класса во время оценки вместо первого.

Я думаю, что я должен работать рядом с этими линиямикода, но я не знаю, как:

test_accuracy, predictions = eval_session.run(
  [evaluation_step, prediction],
  feed_dict={
      bottleneck_input: test_bottlenecks,
      ground_truth_input: test_ground_truth
  })