benchmarks/model_benchmark/benchmark_utils.py/BenchmarkMetricsCallback/on_train_batch_begin
class BenchmarkMetricsCallback:    def on_train_batch_begin(self, batch, logs=None):
        if batch == self.start_batch:
            self.state["epoch_begin_time"] = time.time()