Как получить подробное время выполнения для train_step, используя tf.keras?

Question

как я могу получить подробное время выполнения каждого train_step через tf.keras во время тренировки ?? Хотя я знаю, что добавление обратных вызовов в model.fit () способно возвращать время эпохи, но мне нужно знать: предсказание

  def train_step(inputs):
    images, labels = inputs

    with tf.GradientTape() as tape:
      predictions = model(images, training=True)
      loss = compute_loss(labels, predictions)

    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))

    train_accuracy.update_state(labels, predictions)
    return loss