def _nearest_neighbors_on_grid(
        self,
        input_features: Tensor,
        candidate_features: Tensor,
        input_grid: Tensor,
        candidate_grid: Tensor,
        num_matches: int,
    ) -> Tuple[Tensor, Tensor]:
        """Finds num_matches closest neighbors of input_features in candidate_features
        based on the distance between the features defined by input_grid and
        candidate_grid.

        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.
            input_grid:
                Grid tensor with shape (batch_size, height, width, 2).
            candidate_grid:
                Grid tensor with shape (batch_size, height, width, 2). Note that height
                and width can be different from input_grid.

        Returns:
            (nn_input, nn_candidate) tuple containing two tensors with shape
            (batch_size, num_matches, dim).
        """
        distances: Tensor = torch.cdist(input_grid, candidate_grid)
        # 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
        based on the distance between the features defined by input_grid and
        candidate_grid.

        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.
            input_grid:
                Grid tensor with shape (batch_size, height, width, 2).
            candidate_grid:
                Grid tensor with shape (batch_size, height, width, 2). Note that height
                and width can be different from input_grid.

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