def _local_l2_loss(
        self,
        z_a: Tensor,
        z_b: Tensor,
    ) -> Tensor:
        """Returns loss for local features matched with neareast neighbors using L2
        distance in the feature space.

        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).

        Returns:
            The computed loss for local features.
        """
        # (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)

        # Find nearest neighbours using L2 distance
        z_a_filtered, z_a_nn = self._nearest_neighbors_on_l2(
            input_features=z_a, candidate_features=z_b, num_matches=self.num_matches[0]
        )
        z_b_filtered, z_b_nn = self._nearest_neighbors_on_l2(
            input_features=z_b, candidate_features=z_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 neareast neighbors using L2
        distance in the feature space.

        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).

        Returns:
            The computed loss for local features.
        
lightly/loss/vicregl_loss.py