class_docstr: CrossEntropy loss is calculated in parallel, allreduce denominator into single gpu and calculate softmax.
  Implemented of ArcFace (https://arxiv.org/pdf/1801.07698v1.pdf):
class_docstr_tok: ["CrossEntropy","loss","is","calculated","in","parallel",",","allreduce","denominator","into","single","gpu","and","calculate","softmax",".","Implemented","of","ArcFace","(","https",":","\/\/arxiv.org\/pdf\/1801.07698v1.pdf",")",":"]
code: def forward(ctx, logits: torch.Tensor, label: torch.Tensor):
      """ """
      batch_size = logits.size(0)
      # for numerical stability
      max_logits, _ = torch.max(logits, dim=1, keepdim=True)
      # local to global
      distributed.all_reduce(max_logits, distributed.ReduceOp.MAX)
      logits.sub_(max_logits)
      logits.exp_()
      sum_logits_exp = torch.sum(logits, dim=1, keepdim=True)
      # local to global
      distributed.all_reduce(sum_logits_exp, distributed.ReduceOp.SUM)
      logits.div_(sum_logits_exp)
      index = torch.where(label != -1)[0]
      # loss
      loss = torch.zeros(batch_size, 1, device=logits.device)
      loss[index] = logits[index].gather(1, label[index])
      distributed.all_reduce(loss, distributed.ReduceOp.SUM)
      ctx.save_for_backward(index, logits, label)
      return loss.clamp_min_(1e-30).log_().mean() * (-1)
  
imports: import math
  from typing import Callable
  import torch
  from torch import distributed
  from torch.nn.functional import linear, normalize
label_desc: Use image node_id 1 for calling the DistCrossEntropyFunc obj's underlying member method code with example usage: obj.forward(ctx, logits, label) and returns: unknown