def comparison_mask(a_labels, b_labels):
    'Computes boolean mask for distance comparisons'
    return torch.eq(expand_dims(a_labels, 1), expand_dims(b_labels, 0))
Computes boolean mask for distance comparisons
def comparison_mask(a_labels, b_labels):
    
    return torch.eq(expand_dims(a_labels, 1), expand_dims(b_labels, 0))