def sample(self, total_targets):
        """
        Get sub_weights according to total targets gathered from all GPUs, due to each weights in different
        GPU contains different class centers.
        """
        index_positive = (self.class_start <= total_targets) & (total_targets < self.class_start + self.num_local)
        total_targets[~index_positive] = -1
        total_targets[index_positive] -= self.class_start
        if int(self.sample_rate) != 1:
            positive = torch.unique(total_targets[index_positive], sorted=True)
            if self.num_sample - positive.size(0) >= 0:
                perm = torch.rand(size=[self.num_local], device=self.weight.device)
                perm[positive] = 2.0
                index = torch.topk(perm, k=self.num_sample)[1]
                index = index.sort()[0]
            else:
                index = positive
            self.index = index
            total_targets[index_positive] = torch.searchsorted(index, total_targets[index_positive])
            self.sub_weight = nn.Parameter(self.weight[index])
            self.sub_weight_mom = self.weight_mom[index]

        Get sub_weights according to total targets gathered from all GPUs, due to each weights in different
        GPU contains different class centers.
        
projects/FastFace/fastface/modeling/partial_fc.py