import torch


def pdist(A, B, sqrt=True):
    """ Pairwise Euclidean distance
    >>> import torch; A = torch.randn(2,4); B = torch.randn(15,4); pdist(A, B).size()
    torch.Size([2, 15])
    """
    A_squared = A.pow(2).sum(1).unsqueeze(1)
    B_squared = B.pow(2).sum(1).unsqueeze(0)
    AB = torch.mm(A, B.t()) # A @ B.t()
    A_B_squared = A_squared + B_squared - 2 * AB
    A_B_squared.clamp_(min=1e-16) # min=0 will cause Nan grad
    return torch.sqrt( A_B_squared ) if sqrt else A_B_squared