code: def backward(ctx, grad_loss, grad_nll_loss):
      cuda_module = transformer_cuda_module
      backward_func = (
          cuda_module.cross_entropy_layer_bw_fp16
          if ctx.config.fp16
          else cuda_module.cross_entropy_layer_bw_fp32
      )
  
      assert ctx.config.training
  
      (inputs, targets) = ctx.saved_tensors
  
      targets = targets.to(torch.int32)
      grad_loss = grad_loss.to(torch.float32)
      if ctx.config.fp16:
          inputs = inputs.to(torch.half)
  
      (grad_inputs,) = backward_func(
          ctx.config.layer_id, grad_loss, inputs, targets
      )
  
      return (None, grad_inputs, None)
  
file_path: lightseq/lightseq/training/ops/pytorch/cross_entropy_layer.py
imports: from dataclasses import dataclass
  import torch
  from torch import nn
  from torch.autograd import Function
  from lightseq.training.ops.pytorch.builder import TransformerBuilder
label_desc: Use image node_id 2 for calling the LSCrossEntropyFunc obj's underlying member method code with example usage: obj.backward(ctx, grad_loss, grad_nll_loss) and returns: None,  grad_inputs,  None