def sample(self, features: Tensor, labels: TLabels) -> TTriplets:
        """
        This method samples the hardest triplets in the batch.

        Args:
            features: tensor of shape (batch_size; embed_dim) that contains
                k samples for each of p classes
            labels: labels of the batch, list or tensor of size (batch_size)

        Returns:
            p triplets of (mean_vector, positive, negative_mean_vector)
        """
        # Convert labels to list
        labels = convert_labels2list(labels)
        self._check_input_labels(labels)

        # Get matrix of indices of labels in batch
        labels_mask = self._get_labels_mask(labels)
        p = labels_mask.shape[0]

        embed_dim = features.shape[-1]
        # Reshape embeddings to groups of (p, k, embed_dim) ones,
        # each i-th group contains embeddings of i-th class.
        features = features.repeat((p, 1, 1))
        features = features[labels_mask].view((p, -1, embed_dim))

        # Count mean vectors for each class in batch
        mean_vectors = features.mean(1)

        d_intra = self._count_intra_class_distances(features, mean_vectors)
        # Count the distances to the sample farthest from mean vector
        # for each class.
        pos_indices = d_intra.max(1).indices
        # Count matrix of distances from mean vectors to each other
        d_inter = self._count_inter_class_distances(mean_vectors)
        # For each class mean vector get the closest mean vector
        d_inter = self._fill_diagonal(d_inter, float("inf"))
        neg_indices = d_inter.min(1).indices
        positives = torch.stack(
            [features[idx][pos_idx] for idx, pos_idx in enumerate(pos_indices)]
        )

        return mean_vectors, positives, mean_vectors[neg_indices]

        This method samples the hardest triplets in the batch.

        Args:
            features: tensor of shape (batch_size; embed_dim) that contains
                k samples for each of p classes
            labels: labels of the batch, list or tensor of size (batch_size)

        Returns:
            p triplets of (mean_vector, positive, negative_mean_vector)
        
catalyst/contrib/data/sampler_inbatch.py