def METHOD_NAME(self, output: torch.Tensor) -> None:
    probabilities = self._extract_features(output)
    prob_sum = torch.sum(probabilities, 0, dtype=torch.float64)
    log_prob = torch.log(probabilities + self._eps)
    if log_prob.dtype != probabilities.dtype:
        log_prob = log_prob.to(probabilities)
    kl_sum = torch.sum(probabilities * log_prob, 0, dtype=torch.float64)
    self._num_examples += probabilities.shape[0]
    self._prob_total += prob_sum
    self._total_kl_d += kl_sum