def get_updated_group_features(self, x: Tensor) -> Tensor:
        """Performs the synchronous momentum update of the group vectors.

        Args:
            x:
                Tensor of shape bsz x dim.

        Returns:
            The updated group features.
        """
        assignments = self.assign_groups(x)
        group_features = torch.clone(self.group_features.data)
        for assigned_class in torch.unique(assignments):
            mask = assignments == assigned_class
            group_features[assigned_class] = self.beta * self.group_features[
                assigned_class
            ] + (1 - self.beta) * x[mask].mean(dim=0)

        return group_features
Performs the synchronous momentum update of the group vectors.

        Args:
            x:
                Tensor of shape bsz x dim.

        Returns:
            The updated group features.
        
lightly/models/modules/heads.py