code: def get_label(self, logits, num_labels=-1):
      th_logit = torch.zeros_like(logits[..., :1])
      output = torch.zeros_like(logits).to(logits)
      mask = logits > th_logit
      if num_labels > 0:
          top_v, _ = torch.topk(logits, num_labels, dim=1)
          top_v = top_v[:, -1]
          mask = (logits >= top_v.unsqueeze(1)) & mask
      output[mask] = 1.0
      output[:, 0] = (output[:, 1:].sum(1) == 0.0).to(logits)
  
      return output
  
file_path: DeepKE/src/deepke/relation_extraction/document/losses.py
imports: import torch
  import torch.nn
  import torch.nn.functional
label_desc: Use image node_id 3 for calling the ATLoss obj's underlying member method code with example usage: obj.get_label(logits, num_labels) and returns: output