    def get_metric(self, reset: bool = False) -> Dict[int, Dict[int, torch.FloatTensor]]:
        """
        # Returns

        distances : `Dict[int, Dict[int, torch.FloatTensor]]`
            A dictionary mapping each class label y to a dictionary mapping each protected
            variable label a to the KL divergence or Wasserstein distance of P(C | A = a, Y = y) from
            P(C | Y = y). A distance of nearly 0 implies fairness on the basis of Separation.

        !!! Note
            If a class label is not present in Y conditioned on a protected variable label,
            the expected behavior is that the KL divergence corresponding to this (class label, protected variable
            label) pair is NaN. You can avoid this by using Wasserstein distance instead.
        """
        distances: Dict[int, Dict[int, torch.FloatTensor]] = {}
        if self._total_predictions == 0:
            distances = {
                y: {
                    a: torch.tensor(float("nan"))
                    for a in range(self._num_protected_variable_labels)
                }
                for y in range(self._num_classes)
            }
            return distances

        for y in range(self._num_classes):
            probs = self._predicted_label_counts_by_gold_label[y] / self._total_predictions
            C_given_y_dist = Categorical(probs)
            if self._dist_metric == wasserstein_distance:
                C_given_y_dist = C_given_y_dist.probs
            distances[y] = {}
            for a in range(self._num_protected_variable_labels):
                probs = (
                    self._predicted_label_counts_by_gold_label_and_protected_variable_label[y][a]
                    / self._total_predictions
                )
                if self._dist_metric == kl_divergence:
                    # Implies class label y is not present in Y conditioned on protected variable label a
                    if probs.sum() == 0:
                        distances[y][a] = torch.tensor(float("nan"))
                        continue
                    C_given_a_and_y_dist = Categorical(probs)
                    distances[y][a] = self._dist_metric(C_given_a_and_y_dist, C_given_y_dist)
                elif self._dist_metric == wasserstein_distance:
                    C_given_a_and_y_dist = Categorical(probs).probs
                    # scipy expects empirical values as well as probabilities
                    label_values = torch.tensor(range(self._num_classes))
                    distances[y][a] = self._dist_metric(
                        label_values, label_values, C_given_a_and_y_dist, C_given_y_dist
                    )
        if reset:
            self.reset()
        return distances
# Returns

distances : `Dict[int, Dict[int, torch.FloatTensor]]`
A dictionary mapping each class label y to a dictionary mapping each protected
variable label a to the KL divergence or Wasserstein distance of P(C | A = a, Y = y) from
P(C | Y = y). A distance of nearly 0 implies fairness on the basis of Separation.

!!! Note
If a class label is not present in Y conditioned on a protected variable label,
the expected behavior is that the KL divergence corresponding to this (class label, protected variable
label) pair is NaN. You can avoid this by using Wasserstein distance instead.
allennlp/fairness/fairness_metrics.py