def METHOD_NAME(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)