import torch


def pairwise_distance(point_cloud):
    """
    Args:
    point_cloud: tensor (batch_size, num_points, num_dims)

    Returns:
    pairwise distance: (batch_size, num_points, num_points)
    """
    batch_size = point_cloud.size()[0]
    point_cloud = torch.squeeze(point_cloud)
    if batch_size==1:
        point_cloud = point_cloud.unsqueeze(0)
    point_cloud_transpose = point_cloud.permute(0, 2, 1)
    point_cloud_inner = -2*torch.bmm(point_cloud, point_cloud_transpose)
    point_cloud_square = (point_cloud**2).sum(dim=-1, keepdim=True)
    point_cloud_square_transpose = point_cloud_square.permute(0, 2, 1)
    return point_cloud_square + point_cloud_inner + point_cloud_square_transpose