code: def backward(self, out_grads=None):
      # print('in backward')
      assert self.binded and self.params_initialized
      # tmp_ctx = self._ctx_cpu
      tmp_ctx = self._ctx_single_gpu
      fc7_outs = []
      ctx_fc7_max = self.get_ndarray(
          tmp_ctx, "ctx_fc7_max", (self._batch_size, len(self._context))
      )
      # local_fc7_max = nd.zeros( (self.global_label.shape[0],1), ctx=mx.cpu())
      for i, _module in enumerate(self._arcface_modules):
          _fc7 = _module.get_outputs(merge_multi_context=True)[0]
          fc7_outs.append(_fc7)
          _fc7_max = nd.max(_fc7, axis=1).as_in_context(tmp_ctx)
          ctx_fc7_max[:, i] = _fc7_max
  
      local_fc7_max = self.get_ndarray(
          tmp_ctx, "local_fc7_max", (self._batch_size, 1)
      )
      nd.max(ctx_fc7_max, axis=1, keepdims=True, out=local_fc7_max)
      global_fc7_max = local_fc7_max
      # local_fc7_sum = None
      local_fc7_sum = self.get_ndarray(
          tmp_ctx, "local_fc7_sum", (self._batch_size, 1)
      )
      local_fc7_sum[:, :] = 0.0
      for i, _module in enumerate(self._arcface_modules):
          _max = self.get_ndarray2(
              fc7_outs[i].context, "fc7_max", global_fc7_max
          )
          fc7_outs[i] = nd.broadcast_sub(fc7_outs[i], _max)
          fc7_outs[i] = nd.exp(fc7_outs[i])
          _sum = nd.sum(
              fc7_outs[i], axis=1, keepdims=True
          ).as_in_context(tmp_ctx)
          local_fc7_sum += _sum
      global_fc7_sum = local_fc7_sum
  
      if self._iter % self._verbose == 0:
          # _ctx = self._context[-1]
          _ctx = self._ctx_cpu
          _probs = []
          for i, _module in enumerate(self._arcface_modules):
              _prob = self.get_ndarray2(
                  _ctx, "_fc7_prob_%d" % i, fc7_outs[i]
              )
              _probs.append(_prob)
          fc7_prob = self.get_ndarray(
              _ctx,
              "test_fc7_prob",
              (
                  self._batch_size,
                  self._ctx_num_classes * len(self._context),
              ),
          )
          nd.concat(*_probs, dim=1, out=fc7_prob)
          fc7_pred = nd.argmax(fc7_prob, axis=1)
          local_label = self.global_label - self._local_class_start
          # local_label = self.get_ndarray2(_ctx, 'test_label', local_label)
          _pred = nd.equal(fc7_pred, local_label)
          print("{fc7_acc}", self._iter, nd.mean(_pred).asnumpy()[0])
  
      # local_fc1_grad = []
      # fc1_grad_ctx = self._ctx_cpu
      fc1_grad_ctx = self._ctx_single_gpu
      local_fc1_grad = self.get_ndarray(
          fc1_grad_ctx,
          "local_fc1_grad",
          (self._batch_size, self._emb_size),
      )
      local_fc1_grad[:, :] = 0.0
  
      for i, _module in enumerate(self._arcface_modules):
          _sum = self.get_ndarray2(
              fc7_outs[i].context, "fc7_sum", global_fc7_sum
          )
          fc7_outs[i] = nd.broadcast_div(fc7_outs[i], _sum)
          a = i * self._ctx_num_classes
          b = (i + 1) * self._ctx_num_classes
          _label = self.global_label - self._ctx_class_start[i]
          _label = self.get_ndarray2(
              fc7_outs[i].context, "label", _label
          )
          onehot_label = self.get_ndarray(
              fc7_outs[i].context,
              "label_onehot",
              (self._batch_size, self._ctx_num_classes),
          )
          nd.one_hot(
              _label,
              depth=self._ctx_num_classes,
              on_value=1.0,
              off_value=0.0,
              out=onehot_label,
          )
          fc7_outs[i] -= onehot_label
          _module.backward(out_grads=[fc7_outs[i]])
          # ctx_fc1_grad = _module.get_input_grads()[0].as_in_context(mx.cpu())
          ctx_fc1_grad = self.get_ndarray2(
              fc1_grad_ctx,
              "ctx_fc1_grad_%d" % i,
              _module.get_input_grads()[0],
          )
          local_fc1_grad += ctx_fc1_grad
  
      global_fc1_grad = local_fc1_grad
      self._curr_module.backward(out_grads=[global_fc1_grad])
  
imports: import logging
  import copy
  import time
  import os
  import mxnet
  import numpy
  from mxnet import context
  from mxnet.initializer import Uniform
  from mxnet.module.base_module import BaseModule
  from mxnet.module.module import Module
  from mxnet import metric
  from mxnet.model import BatchEndParam
  from mxnet import io
  import mxnet.ndarray
  from config import config
label_desc: Use image node_id 18 for calling the ParallModule obj's underlying member method code with example usage: obj.backward(out_grads) without return types