benchmarks/layer_benchmark/base_benchmark.py/BenchmarkMetricsCallback/on_train_batch_end
class BenchmarkMetricsCallback:    def on_train_batch_end(self, batch, logs=None):
        if batch == self.stop_batch:
            self.state["benchmark_end"] = time.time()
            throughput = (self.stop_batch - self.start_batch + 1) / (
                self.state["benchmark_end"] - self.state["benchmark_begin"]
            )
            self.state["throughput"] = throughput