def __init__(
        self,
        group_features: Tensor,
        beta: float,
    ):
        """Initializes the SMoGPrototypes module with the specified parameter."""
        super(SMoGPrototypes, self).__init__()
        self.group_features = nn.Parameter(group_features, requires_grad=False)
        self.beta = beta
Initializes the SMoGPrototypes module with the specified parameter.
lightly/models/modules/heads.py