    def log_batch(
        self,
        batch_grad_norm: Optional[float],
        metrics: Dict[str, float],
        batch_group: List[TensorDict],
        param_updates: Optional[Dict[str, torch.Tensor]],
        batch_number: int,
    ) -> None:
        """
        Called every batch to perform all of the logging that is due.
        """
        if batch_number <= 1:  # batch_number is usually 1-indexed
            self._cumulative_batch_group_size = 0
            self.log_inputs(batch_group)

        if self._should_log_this_batch():
            if self._should_log_parameter_statistics:
                self._log_parameter_and_gradient_statistics(batch_grad_norm)

            if self._should_log_learning_rate:
                self._log_learning_rates()

            # Now collect per-batch metrics to log.
            metrics_to_log: Dict[str, float] = {}
            batch_loss_metrics = {"batch_loss", "batch_reg_loss"}
            for key in batch_loss_metrics:
                if key not in metrics:
                    continue
                value = metrics[key]
                metrics_to_log[key] = value
                # Update and add moving average.
                self._batch_loss_moving_sum[key] += value
                self._batch_loss_moving_items[key].append(value)
                if len(self._batch_loss_moving_items[key]) > self._batch_loss_moving_average_count:
                    self._batch_loss_moving_sum[key] -= self._batch_loss_moving_items[key].popleft()
                metrics_to_log[f"{key}_mov_avg"] = self._batch_loss_moving_sum[key] / len(
                    self._batch_loss_moving_items[key]
                )

            for key, value in metrics.items():
                if key in batch_loss_metrics:
                    continue
                key = "batch_" + key
                if key not in metrics_to_log:
                    metrics_to_log[key] = value

            self.log_scalars(
                metrics_to_log,
                log_prefix="train",
            )

        if self._should_log_distributions_this_batch():
            assert param_updates is not None
            self._log_distributions()
            self._log_gradient_updates(param_updates)

        if self._batch_size_interval:
            # We're assuming here that `log_batch` will get called every batch, and only every
            # batch. This is true with our current usage of this code (version 1.0); if that
            # assumption becomes wrong, this code will break.
            batch_group_size = sum(get_batch_size(batch) for batch in batch_group)  # type: ignore
            self._cumulative_batch_group_size += batch_group_size
            if batch_number % self._batch_size_interval == 0:
                average = self._cumulative_batch_group_size / batch_number
                self.log_scalars(
                    {"batch_size": batch_group_size, "mean_batch_size": average}, log_prefix="train"
                )
Called every batch to perform all of the logging that is due.
allennlp/training/callbacks/log_writer.py