code: def get_model_norm(model):
      with torch.no_grad():
          total = 0.0
          for param in model.parameters():
              should_gather = (
                  hasattr(param, "ds_id")
                  and param.ds_status == ZeroParamStatus.NOT_AVAILABLE
              )
              with deepspeed.zero.GatheredParameters(
                  param, enabled=should_gather
              ):
                  total += float(param.float().norm())
  
      return total
  
imports: import torch
  import torch.nn.functional
  import time
  import deepspeed
  from deepspeed.runtime.zero.partition_parameters import ZeroParamStatus
  from deepspeed.accelerator import get_accelerator
  from dschat.utils.utils import print_rank_0
label_desc: Use image node_id 2 for calling a global function with example usage: get_model_norm(model) and returns: total