def _stop_trace(self, batch=None):
        """Logs the trace graph to TensorBoard."""
        if batch is None:
            batch = self._stop_batch
        with self._train_writer.as_default():
            # TODO(b/126388999): Remove step info in the summary name.
            self.summary.trace_export(name="batch_%d" % batch, step=batch)
        self._stop_profiler()
        self._is_tracing = False
Logs the trace graph to TensorBoard.
keras/src/callbacks/tensorboard.py