code: def test_encoder_layer_backward():
      batch_size, seq_len = kt.bs_sl()
      print(f"(batch_size, seq_len): ({batch_size}, {seq_len})")
  
      hidden_size = 1024
  
      hidden_states = kt.rand(
          (batch_size, seq_len, hidden_size)
      ).requires_grad_()
      self_attn_padding_mask = kt.attn_mask(
          batch_size, seq_len, dtype=torch.bool
      )
      loss_data = torch.randn(1, dtype=hidden_states.dtype).sum()
  
      def custom():
          for i in range(NUM_LAYERS):
              custom_enc_layers[i].zero_grad()
          res = hidden_states.clone()
          for i in range(NUM_LAYERS):
              res = custom_enc_layers[i](res, self_attn_padding_mask)
          custom_loss = (res / 1000).sum()
          custom_loss.data.copy_(loss_data)
          custom_loss.backward()
          grad_list = []
          for i in range(NUM_LAYERS - 1, -1, -1):
              """
              attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb,
              inter_w, inter_b, output_w, output_b, ffn_nw, ffn_nb
              """
              grads = split_custom_layer_grad(custom_enc_layers[i])
              grad_list.extend(grads[:11])
              pass
          return grad_list
  
      def baseline():
          for i in range(NUM_LAYERS):
              base_enc_layers[i].zero_grad()
          res = hidden_states.clone()
          for i in range(NUM_LAYERS):
              res = base_enc_layers[i](res, self_attn_padding_mask)
          custom_loss = (res / 1000).sum()
          custom_loss.data.copy_(loss_data)
          custom_loss.backward()
          grad_list = []
          for i in range(NUM_LAYERS - 1, -1, -1):
              """
              attn_qkvw, attn_qkvb, attn_ow, attn_ob, attn_nw, attn_nb,
              inter_w, inter_b, output_w, output_b, ffn_nw, ffn_nb
              """
              grads = split_custom_layer_grad(base_enc_layers[i])
              grad_list.extend(grads[:11])
              pass
          return grad_list
  
      return custom, baseline
  
file_path: lightseq/tests/test_ls_layers_new.py
imports: import random
  import math
  from copy import deepcopy
  from dataclasses import dataclass
  import torch
  import torch.nn
  from tests.util import TestDecorator, get_fairseq_enc_params, get_fairseq_dec_params, max_batch_tokens, max_seq_len, split_custom_layer_grad, copy_grad_from_paras
  from tests import fairseq_layers
  from lightseq.training.ops.pytorch.transformer_encoder_layer_new import LSTransformerEncoderLayerNew
  from lightseq.training.ops.pytorch.transformer_encoder_layer import LSTransformerEncoderLayer
  from lightseq.training.ops.pytorch import layer_cuda_module
  from lightseq.training.cli.fs_modules.ls_fs_transformer_decoder_layer import LSFSTransformerDecoderLayerNew, LSFSTransformerDecoderLayer
label_desc: Use image node_id 5 for calling a global function with example usage: test_encoder_layer_backward() and returns: custom,  baseline, grad_list, grad_list