def forward(
        self, input: torch.Tensor, target: torch.LongTensor = None
    ) -> torch.Tensor:
        """
        Args:
            input: input features,
                expected shapes ``BxF`` where ``B``
                is batch dimension and ``F`` is an
                input feature dimension.
            target: target classes,
                expected shapes ``B`` where
                ``B`` is batch dimension.
                If `None` then will be returned
                projection on centroids.
                Default is `None`.

        Returns:
            tensor (logits) with shapes ``BxC``
            where ``C`` is a number of classes
            (out_features).
        """
        cos_theta = F.linear(F.normalize(input), F.normalize(self.weight))
        theta = torch.acos(torch.clamp(cos_theta, -1.0 + self.eps, 1.0 - self.eps))

        if target is None:
            return cos_theta

        one_hot = torch.zeros_like(cos_theta)
        one_hot.scatter_(1, target.view(-1, 1).long(), 1)

        if self.train:
            with torch.no_grad():
                b_avg = (
                    torch.where(
                        one_hot < 1,
                        torch.exp(self.s * cos_theta),
                        torch.zeros_like(cos_theta),
                    )
                    .sum(1)
                    .mean()
                )
                theta_median = theta[one_hot > 0].median()
                theta_median = torch.min(
                    torch.full_like(theta_median, math.pi / 4), theta_median
                )
                self.s = (torch.log(b_avg) / torch.cos(theta_median)).item()

        logits = self.s * cos_theta
        return logits

        Args:
            input: input features,
                expected shapes ``BxF`` where ``B``
                is batch dimension and ``F`` is an
                input feature dimension.
            target: target classes,
                expected shapes ``B`` where
                ``B`` is batch dimension.
                If `None` then will be returned
                projection on centroids.
                Default is `None`.

        Returns:
            tensor (logits) with shapes ``BxC``
            where ``C`` is a number of classes
            (out_features).
        
catalyst/contrib/layers/cosface.py