When I was using Insightface, I found that sometimes the function sample took up a lot of time, resulting in low GPU utilization. I used line_profiler to check the time of each line, and found that this line of code runs very slowly

  
      github.com
  
  
    deepinsight/insightface/blob/master/recognition/partial_fc/pytorch/partial_classifier.py#L56

        self.sub_weight = Parameter(self.weight)
        self.sub_weight_mom = self.weight_mom
    else:
        self.sub_weight = Parameter(torch.empty((0, 0)).cuda(local_rank))


@torch.no_grad()
def sample(self, total_label):
    P = (self.class_start <=
         total_label) & (total_label < self.class_start + self.num_local)
    total_label[~P] = -1
    total_label[P] -= self.class_start
    if int(self.sample_rate) != 1:
        positive = torch.unique(total_label[P], sorted=True)
        if self.num_sample - positive.size(0) >= 0:
            perm = torch.rand(self.num_local,device=cfg.local_rank)
            perm[positive] = 2.0
            index = torch.topk(perm,k=self.num_sample)[1]
            index = index.sort()[0]
        else:
            index = positive




  
  
    
    
  
  


unreasonable result:
无标题1037×725 29.8 KB
reasonable result (I split this line into two lines):
rBAoL1_J_KWAdMRvAAEu6vsp8vQ6351004×725 29.4 KB
How can I solve this problem? Thanks.