tensorboardのIMAGESで表示させたいが、step(epoch)毎にできない。

やりたいこと

mnistを使って、tenosrboardでいろいろと可視化しました。
tensorboard上でstep数を動かして画像の変化を見たいのですが

file_writer = tf.summary.create_file_writer(log_dir)
with file_writer.as_default():
  images = np.reshape(x_train[0:10], (-1, 28, 28, 1))
  tf.summary.image("25 training data examples", images, max_outputs=10, step=1)

でstepを1など数字で指定するとできるのですが、
step=epochとしても

NameError: name 'epoch' is not defined

と出ます。

↓のようにしたいです。
](479330c82417ff546b0849de47edd1e2.jpeg)

変更点などよろしくお願いします。

コード全体

import tensorflow as tf
import datetime
import numpy as np
mnist = tf.keras.datasets.mnist

(x_train, y_train),(x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(512, activation=tf.nn.relu),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10, activation=tf.nn.softmax)
])
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])



log_dir="logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M")
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1)


model.fit(x=x_train, 
          y=y_train, 
          epochs=5, 
          validation_data=(x_test, y_test), 
          callbacks=[tensorboard_callback]) 


file_writer = tf.summary.create_file_writer(log_dir)
with file_writer.as_default():
  images = np.reshape(x_train[0:10], (-1, 28, 28, 1))
  tf.summary.image("25 training data examples", images, max_outputs=10, step=epoch)


model.evaluate(x_test, y_test)