benchmarks/layer_benchmark/base_benchmark.py/BenchmarkMetricsCallback/on_predict_batch_begin
class BenchmarkMetricsCallback:    def on_predict_batch_begin(self, batch, logs=None):
        if batch == self.start_batch:
            self.state["benchmark_begin"] = time.time()