code: def test_gather_loss():
      # loss could be a Tensor or Dictionary
      loss = torch.tensor([3.2], dtype=torch.float)
  
      stats = Statistics(opts=None, metric_names=["loss"])
      stats.update({}, {}, {"loss": loss})
      np.testing.assert_almost_equal(stats.avg_statistics("loss"), 3.2)
  
      loss_dict = {
          "aux_loss": torch.tensor([4.2]),
          "total_loss": torch.tensor(
              [5.2]
          ),  # total loss key is mandatory for a loss in dict format
      }
      stats = Statistics(opts=None, metric_names=["loss"])
      stats.update({}, {}, {"loss": loss_dict})
      np.testing.assert_almost_equal(
          stats.avg_statistics("loss", "aux_loss"), 4.2, decimal=2
      )
      np.testing.assert_almost_equal(
          stats.avg_statistics("loss", "total_loss"), 5.2, decimal=2
      )
  
      # Empty extras
      stats = Statistics(opts=None, metric_names=["loss"])
      stats.update({}, {}, {})
      stats.update({}, {}, {"loss": None})
  
      metric = stats.avg_statistics("loss")
      np.testing.assert_almost_equal(metric, 0)
  
imports: import numpy
  import torch
  from metrics.stats import Statistics
label_desc: Use image node_id 1 for calling a global function with example usage: test_gather_loss() without return types