code: def get_backward_hook(
      name, trainer, rank, logger, dump_to_file=False
  ):
      """
      A backward hook to dump all of the module input and output grad norms. The hook will be called every time the gradients with respect to module inputs are computed.
      Only float type input/output grad tensor norms are computed.
      For more details about the backward hook, check https://pytorch.org/docs/stable/generated/torch.nn.modules.module.register_module_full_backward_hook.html
  
      Args:
          name: tensor name
          trainer: PTL trainer
          rank: worker rank
          logger: PTL log function
          dump_to_file:  wether dump the csv file to the disk
      """
      if dump_to_file:
          os.makedirs("debug_info", exist_ok=True)
          fp = open(f"debug_info/backward_{name}_rank{rank}.txt", "w")
          header = False
  
      def backward_hook(module, inputs, outputs):
          nonlocal header
          nonlocal fp
          if trainer.training:
              values = []
              headers = []
              for n, i in enumerate(inputs):
                  if isinstance(i, torch.Tensor) and (
                      i.dtype == torch.float
                      or i.dtype == torch.half
                      or i.dtype == torch.bfloat16
                  ):
                      if not header:
                          headers.append("input")
                      input_norm = i.data.norm()
                      values.append(f"{input_norm}")
                      logger(
                          f"debug_info_backward/{name}_rank{rank}_input{n}",
                          input_norm,
                      )
              if isinstance(outputs, tuple):
                  for n, i in enumerate(outputs):
                      if isinstance(i, torch.Tensor) and (
                          i.dtype == torch.float
                          or i.dtype == torch.half
                          or i.dtype == torch.bfloat16
                      ):
                          if not header:
                              headers.append("output")
                          output_norm = i.data.norm()
                          values.append(f"{output_norm}")
                          logger(
                              f"debug_info_backward/{name}_rank{rank}_output{n}",
                              output_norm,
                          )
              else:
                  headers.append("output")
                  values.append(f"{outputs.data.norm()}")
              values.append(f"{trainer.global_step}")
              if not header:
                  headers.append("step")
                  fp.write(",".join(headers) + "\n")
                  header = True
              fp.write(",".join(values) + "\n")
          fp.flush()
  
      return backward_hook
  
imports: import os
  import torch
label_desc: Use image node_id 2 for calling a global function with example usage: get_backward_hook(name, trainer, rank, logger, dump_to_file) and returns: backward_hook