code: def get_tensor_hook(
      module, name, trainer, rank, logger, dump_to_file=False
  ):
      """
      A tensor hook to dump all of the tensor weight norms and grad norms at the end of each of the backward steps.
      For more details about the tensor hook, check https://pytorch.org/docs/stable/generated/torch.Tensor.register_hook.html
  
      Args:
          module: the model module
          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/tensor_{name}_rank{rank}.csv", "w")
          header = False
  
      def tensor_hook(grad):
          nonlocal header
          nonlocal fp
          values = []
          headers = []
  
          weight = module.get_parameter(name)
          weight_norm = weight.data.norm()
          grad_norm = grad.data.norm()
          logger(
              f"debug_info_tensors/{name}_rank{rank}_grad_norm",
              grad_norm,
          )
          logger(
              f"debug_info_tensors/{name}_rank{rank}_weight_norm",
              weight_norm,
          )
          values.append(f"{weight_norm}")
          values.append(f"{grad_norm}")
          values.append(f"{trainer.global_step}")
          if dump_to_file:
              if not header:
                  headers.append("weight")
                  headers.append("grad")
                  headers.append("step")
                  fp.write(",".join(headers) + "\n")
                  header = True
              fp.write(",".join(values) + "\n")
              fp.flush()
          return grad
  
      return tensor_hook
  
imports: import os
  import torch
label_desc: Use image node_id 3 for calling a global function with example usage: get_tensor_hook(module, name, trainer, rank, logger, dump_to_file) and returns: tensor_hook, grad