def test_cluster_count_inter_class_distances(mean_vectors, expected) -> None:
    """
    Test _count_inter_class_distances method of HardClusterSampler.

    Args:
        mean_vectors: tensor of shape (p, embed_dim) -- mean vectors of
        classes in the batch
        expected: tensor of shape (p, p) -- expected distances from mean
        vectors of classes
    """
    sampler = HardClusterSampler()
    distances = sampler._count_inter_class_distances(mean_vectors)
    assert (distances == expected).all()

    Test _count_inter_class_distances method of HardClusterSampler.

    Args:
        mean_vectors: tensor of shape (p, embed_dim) -- mean vectors of
        classes in the batch
        expected: tensor of shape (p, p) -- expected distances from mean
        vectors of classes
    
tests/contrib/data/test_sampler_inbatch.py