def initialize_centroid(self, data_loader: DataLoader) -> None:
        """Initialize the centroid of the memory bank.

        Computes the average feature representation of normal samples to
        initialize the memory bank centroids.

        Args:
            data_loader (DataLoader): DataLoader containing normal training
                samples.

        Example:
            >>> from torch.utils.data import DataLoader
            >>> model = CfaModel(...)
            >>> train_loader = DataLoader(...)
            >>> model.initialize_centroid(train_loader)
        """
        device = next(self.feature_extractor.parameters()).device
        with torch.no_grad():
            for i, data in enumerate(tqdm(data_loader)):
                batch = data.image.to(device)
                features = self.feature_extractor(batch)
                features = list(features.values())
                target_features = self.descriptor(features)
                self.memory_bank = ((self.memory_bank * i) + target_features.mean(dim=0, keepdim=True)) / (i + 1)

        self.memory_bank = rearrange(self.memory_bank, "b c h w -> (b h w) c")

        scale = self.get_scale(batch.shape[-2:])

        if self.gamma_c > 1:
            # TODO(samet-akcay): Create PyTorch KMeans class.
            # CVS-122673
            k_means = KMeans(n_clusters=(scale[0] * scale[1]) // self.gamma_c, max_iter=3000)
            cluster_centers = k_means.fit(self.memory_bank.cpu()).cluster_centers_
            self.memory_bank = torch.tensor(cluster_centers, requires_grad=False).to(device)

        self.memory_bank = rearrange(self.memory_bank, "h w -> w h")
Initialize the centroid of the memory bank.

        Computes the average feature representation of normal samples to
        initialize the memory bank centroids.

        Args:
            data_loader (DataLoader): DataLoader containing normal training
                samples.

        Example:
            >>> from torch.utils.data import DataLoader
            >>> model = CfaModel(...)
            >>> train_loader = DataLoader(...)
            >>> model.initialize_centroid(train_loader)
        
src/anomalib/models/image/cfa/torch_model.py