Профиль предсказания вызова Tensorboard

У меня есть обученная сеть, которую я использую для прогнозирования данных изображения. Я хотел бы ускорить предсказание, но я не знаю, как запустить профайлер tensorboard при предсказании. Я предпочитаю не проводить повторное обучение для оптимизации из-за размера модели. Это похоже на предыдущий вопрос, но я Я спрашиваю снова, потому что этот вопрос не получил соответствующих ответов.

Когда я следую этому примеру профилирования , я получаю профиль, но когда я добавляю шаг прогнозирования с тем же обратным вызовом, он не профилирует прогноз.

(previous lines are loading and setting up an MNIST model as in the linked colab)

tboard_callback = tf.keras.callbacks.TensorBoard(log_dir = logs,
                                                  histogram_freq = 1,
                                                  profile_batch='500,520')
model.fit(ds_train,
          epochs=2,
          validation_data=ds_test,
          callbacks = [tboard_callback]
          )
model.predict(ds_test,callbacks=[tboard_callback])

Среда: Windows 10, python 3.8, tensorflow 2.2, tensorboard_plugin_profile 2.4.0