    def compute_distances(self):
        """
        For every centroid m, computes

                          ||M - m[None, :]||_2

        Remarks:
            - We rely on PyTorch's broadcasting to speed up computations
              and reduce the memory overhead
            - Without chunking, the sizes in the broadcasting are modified as:
              (n_centroids x n_samples x out_features) -> (n_centroids x out_features)
            - The broadcasting computation is automatically chunked so that
              the tensors fit into the memory of the GPU
        """

        nb_centroids_chunks = 1

        while True:
            try:
                return torch.cat(
                    [
                        (self.W[None, :, :] - centroids_c[:, :, None]).norm(p=2, dim=1)
                        for centroids_c in self.centroids.chunk(
                            nb_centroids_chunks, dim=0
                        )
                    ],
                    dim=0,
                )
            except RuntimeError:
                nb_centroids_chunks *= 2
For every centroid m, computes

||M - m[None, :]||_2

Remarks:
- We rely on PyTorch's broadcasting to speed up computations
and reduce the memory overhead
- Without chunking, the sizes in the broadcasting are modified as:
(n_centroids x n_samples x out_features) -> (n_centroids x out_features)
- The broadcasting computation is automatically chunked so that
the tensors fit into the memory of the GPU
fairseq/modules/quantization/pq/em.py