def _nearest_neighbors(feature_bank: torch.Tensor, features: torch.Tensor, n_neighbors: int = 1) -> torch.Tensor:
        """Perform the KNN search.

        Args:
            feature_bank (torch.Tensor): Feature bank of shape ``(M, D)`` used for
                KNN search.
            features (torch.Tensor): Input features of shape ``(N, D)``.
            n_neighbors (int, optional): Number of neighbors used in KNN search.
                Defaults to ``1``.

        Returns:
            torch.Tensor: Distances between the input features and their K nearest
                neighbors in the feature bank.
        """
        distances = torch.cdist(features, feature_bank, p=2.0)  # euclidean norm
        if n_neighbors == 1:
            # when n_neighbors is 1, speed up computation by using min instead of topk
            distances, _ = distances.min(1)
            return distances.unsqueeze(1)
        distances, _ = distances.topk(k=n_neighbors, largest=False, dim=1)
        return distances
Perform the KNN search.

        Args:
            feature_bank (torch.Tensor): Feature bank of shape ``(M, D)`` used for
                KNN search.
            features (torch.Tensor): Input features of shape ``(N, D)``.
            n_neighbors (int, optional): Number of neighbors used in KNN search.
                Defaults to ``1``.

        Returns:
            torch.Tensor: Distances between the input features and their K nearest
                neighbors in the feature bank.
        
src/anomalib/models/video/ai_vad/density.py