def _nearest_neighbors_on_l2(
        self, input_features: Tensor, candidate_features: Tensor, num_matches: int
    ) -> Tuple[Tensor, Tensor]:
        """Finds num_matches closest neighbors of input_features in candidate_features.

        Args:
            input_features:
                Local features tensor with shape (batch_size, height * width, dim).
            candidate_features:
                Local features tensor with shape (batch_size, height * width, dim).
                Note that height and width can be different from input_features.

        Returns:
            (nn_input, nn_candidate) tuple containing two tensors with shape
            (batch_size, num_matches, dim).
        """
        distances = torch.cdist(input_features, candidate_features)
        # TODO(Philipp, 12/24): Remove type ignore when utils are typechecked.
        return utils.nearest_neighbors(  # type: ignore[no-any-return]
            input_features, candidate_features, distances, num_matches
        )
Finds num_matches closest neighbors of input_features in candidate_features.

        Args:
            input_features:
                Local features tensor with shape (batch_size, height * width, dim).
            candidate_features:
                Local features tensor with shape (batch_size, height * width, dim).
                Note that height and width can be different from input_features.

        Returns:
            (nn_input, nn_candidate) tuple containing two tensors with shape
            (batch_size, num_matches, dim).
        
lightly/loss/vicregl_loss.py