import torch


def pdist2(feature1, feature2):
    """Compute pairwise (square) distances of features.

    Args:
        feature1 (torch.Tensor): (num_inst1, channels)
        feature2 (torch.Tensor): (num_inst2, channels)

    Returns:
        distance (torch.Tensor): (num_inst1, num_inst2)

    """
    square_sum1 = torch.sum(feature1 ** 2, 1, keepdim=True)
    square_sum2 = torch.sum(feature2 ** 2, 1, keepdim=True)
    square_sum = square_sum1 + square_sum2.transpose(0, 1)
    distance = torch.addmm(square_sum, feature1, feature2.transpose(0, 1), alpha=-2.0)
    return distance