code: def gather_metrics(
      self,
      prediction: Union[Tensor, Dict],
      target: Union[Tensor, Dict],
      extras: Dict[str, Any],
  ) -> Union[Tensor, Dict[str, Tensor]]:
      if extras is None:
          extras = {}
  
      grad_norm = extras.get("grad_norm", None)
  
      if grad_norm is None:
          grad_norm = 0.0
  
      if isinstance(grad_norm, Tensor):
          return grad_norm
      elif isinstance(grad_norm, Number):
          return torch.tensor(grad_norm, device=self.device)
      elif isinstance(grad_norm, Dict):
          grad_norm.pop(None, None)
  
          for k, v in grad_norm.items():
              if isinstance(v, Number):
                  grad_norm[k] = torch.tensor(
                      grad_norm, device=self.device
                  )
              elif isinstance(v, str):
                  del grad_norm[k]
              elif not isinstance(v, Tensor):
                  logger.error(
                      "Grad-norm metric supports Number, Tensor, or Dict of Tensors."
                      f" Got {v} with {type(v)} type under key {k}."
                  )
  
          return grad_norm
      else:
          logger.error(
              "Grad-norm metric supports Number, Tensor, or Dict of Tensors."
              f" Got {grad_norm} with {type(grad_norm)} type."
          )
  
file_path: ml-cvnets/metrics/misc.py
imports: from numbers import Number
  from typing import Any, Dict, Union
  import torch
  from torch import Tensor
  from metrics import METRICS_REGISTRY
  from metrics.metric_base import AverageMetric
  from utils import logger
label_desc: Use image node_id 1 for calling the GradNormMetric obj's underlying member method code with example usage: obj.gather_metrics(prediction, target, extras) and returns: grad_norm, torch, grad_norm