import torch


def kl_div(p, q):
    """
    Computes the KL divergence between two distributions p and q

    Argument:
        p - probability tensor
        q - probability tensor
    """
    return torch.matmul(p, torch.log(p / q).t()).mean()