def compute(self) -> dict[str, torch.Tensor]:
        """Aggregate the state over all processes to compute the metric.

        Returns:
            loss: The loss averaged across all batches as a :class:`~torch.Tensor`.
        """
        # Return average loss over entire dataset
        sum_perplexity = torch.where(
            self.sum_length != 0,
            self.sum_perplexity,
            -1,
        )
        sum_loss = torch.where(self.sum_length != 0, self.sum_loss, -1)
        sum_length = torch.where(self.sum_length != 0, self.sum_length, 1)

        sum_perplexity_seq_id = torch.where(
            self.sum_length_seq_id != 0,
            self.sum_perplexity_seq_id,
            -1,
        )
        sum_loss_seq_id = torch.where(
            self.sum_length_seq_id != 0,
            self.sum_loss_seq_id,
            -1,
        )
        sum_length_seq_id = torch.where(
            self.sum_length_seq_id != 0,
            self.sum_length_seq_id,
            1,
        )

        return {
            'mean_loss_v_len':
                sum_loss / sum_length,
            'mean_perplexity_v_len':
                sum_perplexity / sum_length,
            'sum_length':
                self.sum_length,
            'mean_loss_seq_id_v_len':
                sum_loss_seq_id / sum_length_seq_id,
            'mean_perplexity_seq_id_v_len':
                sum_perplexity_seq_id / sum_length_seq_id,
            'sum_length_seq_id':
                self.sum_length_seq_id,
        }
Aggregate the state over all processes to compute the metric.

        Returns:
            loss: The loss averaged across all batches as a :class:`~torch.Tensor`.
        
llmfoundry/callbacks/loss_perp_v_len_callback.py