def METHOD_NAME(self, noised_logits, input_logits):
    return (
        F.kl_div(
            F.log_softmax(noised_logits, dim=-1, dtype=torch.float32),
            F.softmax(input_logits, dim=-1, dtype=torch.float32),
            None,
            None,
            "sum",
        )
        + F.kl_div(
            F.log_softmax(input_logits, dim=-1, dtype=torch.float32),
            F.softmax(noised_logits, dim=-1, dtype=torch.float32),
            None,
            None,
            "sum",
        )
    ) / noised_logits.size(0)