def __init__(self, in_features, out_features, sigma=True):
        """
        :param in_features: number of input features
        :param out_features: number of classes
        :param sigma: learnable output scaling factor
        """
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.weight = nn.Parameter(torch.Tensor(self.out_features, in_features))
        if sigma:
            self.sigma = nn.Parameter(torch.Tensor(1))
        else:
            self.register_parameter("sigma", None)
        self.reset_parameters()

        :param in_features: number of input features
        :param out_features: number of classes
        :param sigma: learnable output scaling factor
        
avalanche/models/cosine_layer.py