class_docstr: Modified from Partial-FC
class_docstr_tok: ["Modified","from","Partial-FC"]
code: def prepare(self, label, optimizer):
      with torch.cuda.stream(self.stream):
          total_label = torch.zeros(
              size=[self.batch_size * self.world_size],
              device=self.device,
              dtype=torch.long,
          )
          dist.all_gather(
              list(total_label.chunk(self.world_size, dim=0)), label
          )
          index_positive = self.sample(total_label)
          optimizer.state.pop(
              optimizer.param_groups[-1]["params"][0], None
          )
          optimizer.param_groups[-1]["params"][0] = self.sub_weight
          optimizer.state[self.sub_weight][
              "momentum_buffer"
          ] = self.sub_weight_mom
          norm_weight = normalize(self.sub_weight)
          return total_label, norm_weight, index_positive
  
file_path: insightface/recognition/vpl/vpl.py
imports: import logging
  import os
  import torch
  import torch.distributed
  from torch.nn import Module
  from torch.nn.functional import normalize, linear
  from torch.nn.parameter import Parameter
label_desc: Use image node_id 6 for calling the VPL obj's underlying member method code with example usage: obj.prepare(label, optimizer) and returns: total_label,  norm_weight,  index_positive