def _local_location_loss(
        self,
        z_a: Tensor,
        z_b: Tensor,
        grid_a: Tensor,
        grid_b: Tensor,
    ) -> Tensor:
        """Returns loss for local features matched with nearest neighbors based on
        the feature location.

        Args:
            z_a:
                Local feature tensor with shape (batch_size, height, width, dim).
            z_b:
                Local feature tensor with shape (batch_size, height, width, dim).
                Note that height and width can be different from z_a.
            grid_a:
                Grid tensor with shape (batch_size, height, width, 2).
            grid_b:
                Grid tensor with shape (batch_size, height, width, 2).
                Note that height and width can be different from grid_a.

        Returns:
            The computed loss for local features based on nearest neighbour matching.
        """
        # (batch_size, height, width, dim) -> (batch_size, height * width, dim)
        z_a = z_a.flatten(start_dim=1, end_dim=2)
        z_b = z_b.flatten(start_dim=1, end_dim=2)

        # (batch_size, height, width, 2) -> (batch_size, height * width, 2)
        grid_a = grid_a.flatten(start_dim=1, end_dim=2)
        grid_b = grid_b.flatten(start_dim=1, end_dim=2)

        # Find nearest neighbours based on grid location
        z_a_filtered, z_a_nn = self._nearest_neighbors_on_grid(
            input_features=z_a,
            candidate_features=z_b,
            input_grid=grid_a,
            candidate_grid=grid_b,
            num_matches=self.num_matches[0],
        )
        z_b_filtered, z_b_nn = self._nearest_neighbors_on_grid(
            input_features=z_b,
            candidate_features=z_a,
            input_grid=grid_b,
            candidate_grid=grid_a,
            num_matches=self.num_matches[1],
        )

        # Compute VICReg losses
        loss_a = self.vicreg_loss.forward(z_a=z_a_filtered, z_b=z_a_nn)
        loss_b = self.vicreg_loss.forward(z_a=z_b_filtered, z_b=z_b_nn)
        return 0.5 * (loss_a + loss_b)
Returns loss for local features matched with nearest neighbors based on
        the feature location.

        Args:
            z_a:
                Local feature tensor with shape (batch_size, height, width, dim).
            z_b:
                Local feature tensor with shape (batch_size, height, width, dim).
                Note that height and width can be different from z_a.
            grid_a:
                Grid tensor with shape (batch_size, height, width, 2).
            grid_b:
                Grid tensor with shape (batch_size, height, width, 2).
                Note that height and width can be different from grid_a.

        Returns:
            The computed loss for local features based on nearest neighbour matching.
        
lightly/loss/vicregl_loss.py