code: def unitwise_norm(x, norm_type=2.0):
      if x.ndim <= 1:
          return x.norm(norm_type)
      else:
          # works for nn.ConvNd and nn,Linear where output dim is first in the kernel/weight tensor
          # might need special cases for other weights (possibly MHA) where this may not be true
          return x.norm(
              norm_type, dim=tuple(range(1, x.ndim)), keepdim=True
          )
  
label_desc: Use image node_id 1 for calling a global function with example usage: unitwise_norm(x, norm_type) and returns: x, x