def forward(self, features1: torch.Tensor, features2: torch.Tensor) -> torch.Tensor:
        """

        Args:
            features1: batch with samples features of shape
                [bs; feature_len]
            features2: batch with samples features of shape
                [bs; feature_len]

        Returns:
            torch.Tensor: NTXent loss
        """
        assert (
            features1.shape == features2.shape
        ), f"Invalid shape of input features: {features1.shape} and {features2.shape}"

        feature_matrix = torch.cat([features1, features2])
        feature_matrix = torch.nn.functional.normalize(feature_matrix)
        # if ||x|| = ||y|| = 1 then||x-y||^2 = 2 - 2<x,y>
        cosine_matrix = (2 - torch.cdist(feature_matrix, feature_matrix) ** 2) / 2

        # todo try different places for temparature
        exp_cosine_matrix = torch.exp(cosine_matrix / self.tau)
        # neg part of the loss
        # torch.exp(1) self similarity
        exp_sim_sum = exp_cosine_matrix.sum(dim=1) - e ** (1 / self.tau)
        neg_loss = torch.log(exp_sim_sum)
        pos_loss = self.cosine_sim(features1, features2) / self.tau
        pos_loss = torch.cat([pos_loss, pos_loss])

        # 2*poss_loss (i,j) and (j,i)
        loss = -pos_loss + neg_loss
        if self.reduction == "mean":
            loss = loss.mean()
        elif self.reduction == "sum":
            loss = loss.sum()
        return loss


        Args:
            features1: batch with samples features of shape
                [bs; feature_len]
            features2: batch with samples features of shape
                [bs; feature_len]

        Returns:
            torch.Tensor: NTXent loss
        
catalyst/contrib/losses/ntxent.py