def forward(
        self, x: Tensor, group_features: Tensor, temperature: float = 0.1
    ) -> Tensor:
        """Computes the logits for given model outputs and group features.

        Args:
            x:
                Tensor of shape bsz x dim.
            group_features:
                Momentum updated group features of shape n_groups x dim.
            temperature:
                Temperature parameter for calculating the logits.

        Returns:
            The computed logits.
        """
        x = torch.nn.functional.normalize(x, dim=1)
        group_features = torch.nn.functional.normalize(group_features, dim=1)
        logits = torch.mm(x, group_features.t())
        return logits / temperature
Computes the logits for given model outputs and group features.

        Args:
            x:
                Tensor of shape bsz x dim.
            group_features:
                Momentum updated group features of shape n_groups x dim.
            temperature:
                Temperature parameter for calculating the logits.

        Returns:
            The computed logits.
        
lightly/models/modules/heads.py