    def __call__(
        self,
        predictions: torch.Tensor,
        gold_labels: torch.Tensor,
        mask: Optional[torch.BoolTensor] = None,
    ):
        """
        # Parameters

        predictions : `torch.Tensor`, required.
            A tensor of predictions of shape (batch_size, ...).
        gold_labels : `torch.Tensor`, required.
            A tensor of the same shape as `predictions`.
        mask : `torch.BoolTensor`, optional (default = `None`).
            A tensor of the same shape as `predictions`.
        """
        predictions, gold_labels, mask = self.detach_tensors(predictions, gold_labels, mask)

        # Flatten predictions, gold_labels, and mask. We calculate the covariance between
        # the vectors, since each element in the predictions and gold_labels tensor is assumed
        # to be a separate observation.
        predictions = predictions.view(-1)
        gold_labels = gold_labels.view(-1)

        if mask is not None:
            mask = mask.view(-1)
            predictions = predictions * mask
            gold_labels = gold_labels * mask
            num_batch_items = torch.sum(mask).item()
        else:
            num_batch_items = gold_labels.numel()

        # Note that self._total_count must be a float or int at all times
        # If it is a 1-dimension Tensor, the previous count will equal the updated_count.
        # The sampe applies for previous_total_prediction_mean and
        # previous_total_label_mean below -- we handle this in the code by
        # calling .item() judiciously.
        previous_count = self._total_count
        updated_count = previous_count + num_batch_items

        batch_mean_prediction = torch.sum(predictions) / num_batch_items

        delta_mean_prediction = (
            (batch_mean_prediction - self._total_prediction_mean) * num_batch_items
        ) / updated_count
        previous_total_prediction_mean = self._total_prediction_mean

        batch_mean_label = torch.sum(gold_labels) / num_batch_items
        delta_mean_label = (
            (batch_mean_label - self._total_label_mean) * num_batch_items
        ) / updated_count
        previous_total_label_mean = self._total_label_mean

        batch_coresiduals = (predictions - batch_mean_prediction) * (gold_labels - batch_mean_label)

        if mask is not None:
            batch_co_moment = torch.sum(batch_coresiduals * mask)
        else:
            batch_co_moment = torch.sum(batch_coresiduals)

        delta_co_moment = batch_co_moment + (
            previous_total_prediction_mean - batch_mean_prediction
        ) * (previous_total_label_mean - batch_mean_label) * (
            previous_count * num_batch_items / updated_count
        )

        # Due to the online computation of covariance, the following code can
        # still lead to nan values in later iterations. To be revisited.

        # if is_distributed():

        #     # Note: this gives an approximate aggregation of the covariance.
        #     device = gold_labels.device
        #     delta_mean_prediction = torch.tensor(delta_mean_prediction, device=device)
        #     delta_mean_label = torch.tensor(delta_mean_label, device=device)
        #     delta_co_moment = torch.tensor(delta_co_moment, device=device)
        #     _total_count = torch.tensor(updated_count, device=device)
        #     dist.all_reduce(delta_mean_prediction, op=dist.ReduceOp.SUM)
        #     dist.all_reduce(delta_mean_label, op=dist.ReduceOp.SUM)
        #     dist.all_reduce(delta_co_moment, op=dist.ReduceOp.SUM)
        #     dist.all_reduce(_total_count, op=dist.ReduceOp.SUM)
        #     updated_count = _total_count.item()

        self._total_prediction_mean += delta_mean_prediction.item()
        self._total_label_mean += delta_mean_label.item()
        self._total_co_moment += delta_co_moment.item()
        self._total_count = updated_count
# Parameters

predictions : `torch.Tensor`, required.
A tensor of predictions of shape (batch_size, ...).
gold_labels : `torch.Tensor`, required.
A tensor of the same shape as `predictions`.
mask : `torch.BoolTensor`, optional (default = `None`).
A tensor of the same shape as `predictions`.
allennlp/training/metrics/covariance.py