    def __call__(
        self,
        predicted_labels: torch.Tensor,
        protected_variable_labels: torch.Tensor,
        mask: Optional[torch.BoolTensor] = None,
    ) -> None:
        """
        # Parameters

        predicted_labels : `torch.Tensor`, required.
            A tensor of predicted integer class labels of shape (batch_size, ...). Represented as Y.
        protected_variable_labels : `torch.Tensor`, required.
            A tensor of integer protected variable labels of shape (batch_size, ...). It must be the same
            shape as the `predicted_labels` tensor. Represented as X.
        mask : `torch.BoolTensor`, optional (default = `None`).
            A tensor of the same shape as `predicted_labels`.

        !!! Note
            All tensors are expected to be on the same device.
        """
        predicted_labels, protected_variable_labels, mask = self.detach_tensors(
            predicted_labels, protected_variable_labels, mask
        )

        # Some sanity checks.
        if predicted_labels.size() != protected_variable_labels.size():
            raise ConfigurationError(
                "protected_variable_labels must be of same size as predicted_labels but "
                "found tensor of shape: {}".format(protected_variable_labels.size())
            )
        if mask is not None and predicted_labels.size() != mask.size():
            raise ConfigurationError(
                "mask must be of same size as predicted_labels but "
                "found tensor of shape: {}".format(mask.size())
            )
        if (predicted_labels >= self._num_classes).any():
            raise ConfigurationError(
                "predicted_labels contains an id >= {}, "
                "the number of classes.".format(self._num_classes)
            )
        if (protected_variable_labels >= self._num_protected_variable_labels).any():
            raise ConfigurationError(
                "protected_variable_labels contains an id >= {}, "
                "the number of protected variable labels.".format(
                    self._num_protected_variable_labels
                )
            )

        device = predicted_labels.device
        self._joint_counts_by_protected_variable_label = (
            self._joint_counts_by_protected_variable_label.to(device)
        )
        self._protected_variable_label_counts = self._protected_variable_label_counts.to(device)
        self._y_counts = self._y_counts.to(device)
        self._total_predictions = self._total_predictions.to(device)

        if mask is not None:
            predicted_labels = predicted_labels[mask]
            protected_variable_labels = protected_variable_labels[mask]
        else:
            predicted_labels = predicted_labels.flatten()
            protected_variable_labels = protected_variable_labels.flatten()

        _total_predictions = torch.tensor(predicted_labels.nelement()).to(device)
        _y_counts = torch.zeros(self._num_classes).to(device)
        _y_counts = torch.zeros_like(_y_counts, dtype=predicted_labels.dtype).scatter_add_(
            0, predicted_labels, torch.ones_like(predicted_labels)
        )

        _joint_counts_by_protected_variable_label = torch.zeros(
            (self._num_protected_variable_labels, self._num_classes)
        ).to(device)
        _protected_variable_label_counts = torch.zeros(self._num_protected_variable_labels).to(
            device
        )
        for x in range(self._num_protected_variable_labels):
            x_mask = (protected_variable_labels == x).long()

            _joint_counts_by_protected_variable_label[x] = torch.zeros(self._num_classes).to(device)
            _joint_counts_by_protected_variable_label[x] = torch.zeros_like(
                _joint_counts_by_protected_variable_label[x], dtype=x_mask.dtype
            ).scatter_add_(0, predicted_labels, x_mask)

            _protected_variable_label_counts[x] = torch.tensor(x_mask.sum()).to(device)

        if is_distributed():
            _total_predictions = _total_predictions.to(device)
            dist.all_reduce(_total_predictions, op=dist.ReduceOp.SUM)

            _y_counts = _y_counts.to(device)
            dist.all_reduce(_y_counts, op=dist.ReduceOp.SUM)

            _joint_counts_by_protected_variable_label = (
                _joint_counts_by_protected_variable_label.to(device)
            )
            dist.all_reduce(_joint_counts_by_protected_variable_label, op=dist.ReduceOp.SUM)

            _protected_variable_label_counts = _protected_variable_label_counts.to(device)
            dist.all_reduce(_protected_variable_label_counts, op=dist.ReduceOp.SUM)

        self._total_predictions += _total_predictions
        self._y_counts += _y_counts
        self._joint_counts_by_protected_variable_label += _joint_counts_by_protected_variable_label
        self._protected_variable_label_counts += _protected_variable_label_counts
# Parameters

predicted_labels : `torch.Tensor`, required.
A tensor of predicted integer class labels of shape (batch_size, ...). Represented as Y.
protected_variable_labels : `torch.Tensor`, required.
A tensor of integer protected variable labels of shape (batch_size, ...). It must be the same
shape as the `predicted_labels` tensor. Represented as X.
mask : `torch.BoolTensor`, optional (default = `None`).
A tensor of the same shape as `predicted_labels`.

!!! Note
All tensors are expected to be on the same device.
allennlp/fairness/bias_metrics.py