441 class DistCrossEntropyFunc(torch.autograd.Function):
    442     """
    443     CrossEntropy loss is calculated in parallel, allreduce denominator into single gpu and calculate softmax.
    444     Implemented of ArcFace (https://arxiv.org/pdf/1801.07698v1.pdf):
    445     """
    446 
    447     @staticmethod
    448     def forward(ctx, logits: torch.Tensor, label: torch.Tensor):
    449         """ """
    450         batch_size = logits.size(0)
    451         # for numerical stability
    452         max_logits, _ = torch.max(logits, dim=1, keepdim=True)
    453         # local to global
    454         distributed.all_reduce(max_logits, distributed.ReduceOp.MAX)
    455         logits.sub_(max_logits)
    456         logits.exp_()
    457         sum_logits_exp = torch.sum(logits, dim=1, keepdim=True)
    458         # local to global
    459         distributed.all_reduce(sum_logits_exp, distributed.ReduceOp.SUM)
    460         logits.div_(sum_logits_exp)
    461         index = torch.where(label != -1)[0]
    462         # loss
    463         loss = torch.zeros(batch_size, 1, device=logits.device)
    464         loss[index] = logits[index].gather(1, label[index])
    465         distributed.all_reduce(loss, distributed.ReduceOp.SUM)
    466         ctx.save_for_backward(index, logits, label)
    467         return loss.clamp_min_(1e-30).log_().mean() * (-1)

IndexPut
Traceback (most recent call last):
  File "train.py", line 214, in <module>
    main(parser.parse_args())
  File "train.py", line 157, in main
    loss: torch.Tensor = module_partial_fc(local_embeddings, local_labels, opt)
  File "/home/ma-user/anaconda3/envs/Pytorch-1.11.0/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/cache/development-space/arcface_torch/partial_fc.py", line 213, in forward
    loss = self.dist_cross_entropy(logits, labels)
  File "/home/ma-user/anaconda3/envs/Pytorch-1.11.0/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1110, in _call_impl
    return forward_call(*input, **kwargs)
  File "/cache/development-space/arcface_torch/partial_fc.py", line 498, in forward
    return DistCrossEntropyFunc.apply(logit_part, label_part)
  File "/cache/development-space/arcface_torch/partial_fc.py", line 464, in forward
    loss[index] = logits[index].gather(1, label[index])
RuntimeError: Run:/usr1/workspace/FPTA_Daily_open_pytorchv1.11.0-3.0.tr6/CODE/torch_npu/csrc/framework/OpParamMaker.cpp:136 NPU error,NPU error code is:500002
EZ9999: Inner Error, Please contact support engineer!
EZ9999  The input dtype of x1 x2 y is equal, please check![FUNC:IndexPutVerify][FILE:matrix_calculation_ops.cc][LINE:4676]
        TraceBack (most recent call last):
        Verifying IndexPut failed.[FUNC:InferShapeAndType][FILE:infershape_pass.cc][LINE:135]
        Call InferShapeAndType for node:IndexPut(IndexPut) failed[FUNC:Infer][FILE:infershape_pass.cc][LINE:117]
        process pass InferShapePass on node:IndexPut failed, ret:4294967295[FUNC:RunPassesOnNode][FILE:base_pass.cc][LINE:530]
        build graph failed, graph id:140, ret:1343242270[FUNC:BuildModel][FILE:ge_generator.cc][LINE:1484]
        [Build][SingleOpModel]call ge interface generator.BuildSingleOpModel failed. ge result = 1343242270[FUNC:ReportCallError][FILE:log_inner.cpp][LINE:161]
        [Build][Op]Fail to build op model[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]
        build op model failed, result = 500002[FUNC:ReportInnerError][FILE:log_inner.cpp][LINE:145]
        build graph failed, graph id:141, ret:1343242270[FUNC:BuildModel][FILE:ge_generator.cc][LINE:1484]
        build graph failed, graph id:142, ret:1343242270[FUNC:BuildModel][FILE:ge_generator.cc][LINE:1484]