    def assign(self):
        """
        Assigns each column of W to its closest centroid, thus essentially
        performing the E-step in train().

        Remarks:
            - The function must be called after train() or after loading
              centroids using self.load(), otherwise it will return empty tensors
        """

        distances = self.compute_distances()  # (n_centroids x out_features)
        self.assignments = torch.argmin(distances, dim=0)  # (out_features)
Assigns each column of W to its closest centroid, thus essentially
performing the E-step in train().

Remarks:
- The function must be called after train() or after loading
centroids using self.load(), otherwise it will return empty tensors
fairseq/modules/quantization/pq/em.py