def forward(self, features: Tensor, piece_ids: Tensor) -> Tensor:
        '''
            input: 
                features [NxD]: the feature maps 
                piece_ids [N]: the group ids for feature maps, positive pairs should have the same id 
        '''
        triplet_loss = 0 
        # ids = torch.unique(piece_ids)
        '''
            anchor order is ordered by index:
            piece_ids[positive[i]] == piece_ids[anchor[i]]
            piece_ids[negative[i]] != piece_ids[anchor[i]]
        '''
        feat_num = piece_ids.shape[0]
        anchor = torch.arange(feat_num)
        positive = [random.sample(set(torch.where(piece_ids == piece_ids[i])[0]), 1)[0] for i in range(feat_num)]
        negative = [random.sample(set(torch.where(piece_ids != piece_ids[i])[0]), 1)[0] for i in range(feat_num)]
        positive = torch.Tensor(positive).long()
        negative = torch.Tensor(negative).long()
        anchor_feats = features[anchor]
        positive_feats = features[positive]
        negative_feats = features[negative]
        # NOTE: in fact, the loss has been reduced by "mean" already, but still very large, so I divide it by feat_num AGAIN
        return F.triplet_margin_loss(anchor_feats, positive_feats, negative_feats,
                margin=self.margin, p=self.p, eps=self.eps, swap=self.swap, reduction=self.reduction) / feat_num
        # return F.triplet_margin_loss(anchor, positive, negative, margin=self.margin, p=self.p,
        #                              eps=self.eps, swap=self.swap, reduction=self.reduction)

            input: 
                features [NxD]: the feature maps 
                piece_ids [N]: the group ids for feature maps, positive pairs should have the same id 
        
utils/triplet.py