benchmarks/model_benchmark/benchmark_utils.py/BenchmarkMetricsCallback/on_train_batch_end
class BenchmarkMetricsCallback:    def on_train_batch_end(self, batch, logs=None):
        if batch == self.stop_batch:
            epoch_end_time = time.time()
            throughput = (self.stop_batch - self.start_batch + 1) / (
                epoch_end_time - self.state["epoch_begin_time"]
            )
            self.state["throughput"].append(throughput)