code: def test_gather_grad_norm():
      # Grad norm could be a Tensor or Dictionary
      grad_norm = torch.tensor([3.2], dtype=torch.float)
  
      stats = Statistics(opts=None, metric_names=["grad_norm"])
      stats.update({}, {}, {"grad_norm": grad_norm})
      np.testing.assert_almost_equal(
          stats.avg_statistics("grad_norm"), 3.2
      )
  
      grad_norm_dict = {
          "dummy_norm_a": torch.tensor([4.2]),
          "dummy_norm_b": torch.tensor([5.2]),
      }
      stats = Statistics(opts=None, metric_names=["grad_norm"])
      stats.update({}, {}, {"grad_norm": grad_norm_dict})
  
      _ = stats.avg_statistics("grad_norm")
  
      np.testing.assert_almost_equal(
          stats.avg_statistics("grad_norm", "dummy_norm_a"),
          4.2,
          decimal=2,
      )
      np.testing.assert_almost_equal(
          stats.avg_statistics("grad_norm", "dummy_norm_b"),
          5.2,
          decimal=2,
      )
  
      # Empty extras
      stats = Statistics(opts=None, metric_names=["grad_norm"])
      stats.update({}, {}, {})
      stats.update({}, {}, {"grad_norm": None})
  
      metric = stats.avg_statistics("grad_norm")
      np.testing.assert_almost_equal(metric, 0)
  
imports: import numpy
  import torch
  from metrics.stats import Statistics
label_desc: Use image node_id 2 for calling a global function with example usage: test_gather_grad_norm() without return types