import torch


def nn_distance(pc1, pc2):
    """
    Compute distance
    Args:
        pc1: (B, N, C)
        pc2: (B, M, C)
    Returns:
        dist: (B, N, M)
    """
    N = pc1.shape[1]
    M = pc2.shape[1]
    pc1_expand_tile = pc1.unsqueeze(2).repeat(1,1,M,1)  # (B, N, M, C)
    pc2_expand_tile = pc2.unsqueeze(1).repeat(1,N,1,1)  # (B, N, M, C)
    pc_diff = pc1_expand_tile - pc2_expand_tile  # (B, N, M, C)

    pc_dist = torch.sqrt(torch.sum(pc_diff**2, dim=-1))  # (B, N, M)
    return pc_dist