def compute_distance(self, target_oriented_features: torch.Tensor) -> torch.Tensor:
        """Compute distances between features and memory bank centroids.

        Args:
            target_oriented_features (torch.Tensor): Features from the descriptor
                network.

        Returns:
            torch.Tensor: Distance tensor.

        Example:
            >>> model = CfaModel(...)
            >>> features = torch.randn(32, 256, 32, 32)  # B x C x H x W
            >>> distances = model.compute_distance(features)
        """
        if target_oriented_features.ndim == 4:
            target_oriented_features = rearrange(target_oriented_features, "b c h w -> b (h w) c")

        features = target_oriented_features.pow(2).sum(dim=2, keepdim=True)
        centers = self.memory_bank.pow(2).sum(dim=0, keepdim=True).to(features.device)
        f_c = 2 * torch.matmul(target_oriented_features, (self.memory_bank.to(features.device)))
        return features + centers - f_c
Compute distances between features and memory bank centroids.

        Args:
            target_oriented_features (torch.Tensor): Features from the descriptor
                network.

        Returns:
            torch.Tensor: Distance tensor.

        Example:
            >>> model = CfaModel(...)
            >>> features = torch.randn(32, 256, 32, 32)  # B x C x H x W
            >>> distances = model.compute_distance(features)
        
src/anomalib/models/image/cfa/torch_model.py