import torch


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

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

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

    Notes:
        This method returns square distances, and is optimized for lower memory and faster speed.

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