code: def _loss_class_test(
      rank: int,
      worldsize: int,
      loss_sum_or_avg: Optional[torch.Tensor],
      num_measurements: Optional[torch.Tensor],
      dist_sync_on_step: bool,
      take_avg_loss: bool,
      check_dist_sync_on_step: bool = True,
      check_batch: bool = True,
      atol: float = 1e-8,
  ):
      """Utility function doing the actual comparison between lightning class metric
      and reference metric.
      Args:
          rank: rank of current process
          worldsize: number of processes
          loss_sum_or_avg: a one dimensional float torch tensor with loss sums or means.
          num_measurements: a one dimensional integer torch tensor with number of values on which sums or means from
              ``loss_sum_or_avg`` were computed.
          dist_sync_on_step: bool, if true will synchronize metric state across processes at each call of the
              method :meth:`forward()`
          take_avg_loss: dict with additional arguments used for class initialization
          check_dist_sync_on_step: bool, if true will check if the metric is also correctly
              calculated per batch per device (and not just at the end)
          check_batch: bool, if true will check if the metric is also correctly
              calculated across devices for each batch (and not just at the end)
      """
      # Instantiate lightning metric
      loss_metric = GlobalAverageLossMetric(
          dist_sync_on_step=dist_sync_on_step,
          take_avg_loss=take_avg_loss,
      )
  
      # verify loss works after being loaded from pickled state
      pickled_metric = pickle.dumps(loss_metric)
      loss_metric = pickle.loads(pickled_metric)
      for i in range(rank, NUM_BATCHES, worldsize):
          batch_result = loss_metric(
              loss_sum_or_avg[i], num_measurements[i]
          )
          if loss_metric.dist_sync_on_step:
              if rank == 0:
                  ddp_loss_sum_or_avg = torch.stack(
                      [loss_sum_or_avg[i + r] for r in range(worldsize)]
                  )
                  ddp_num_measurements = torch.stack(
                      [
                          num_measurements[i + r]
                          for r in range(worldsize)
                      ]
                  )
                  sk_batch_result = reference_loss_func(
                      ddp_loss_sum_or_avg,
                      ddp_num_measurements,
                      take_avg_loss,
                  )
                  # assert for dist_sync_on_step
                  if check_dist_sync_on_step:
                      if sk_batch_result.isnan():
                          assert batch_result.isnan()
                      else:
                          assert np.allclose(
                              batch_result.numpy(),
                              sk_batch_result,
                              atol=atol,
                          ), f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}, i = {i}"
          else:
              ls = loss_sum_or_avg[i : i + 1]
              nm = num_measurements[i : i + 1]
              sk_batch_result = reference_loss_func(
                  ls, nm, take_avg_loss
              )
              # assert for batch
              if check_batch:
                  if sk_batch_result.isnan():
                      assert batch_result.isnan()
                  else:
                      assert np.allclose(
                          batch_result.numpy(),
                          sk_batch_result,
                          atol=atol,
                      ), f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}, i = {i}"
      # check on all batches on all ranks
      result = loss_metric.compute()
      assert isinstance(result, torch.Tensor)
      sk_result = reference_loss_func(
          loss_sum_or_avg, num_measurements, take_avg_loss
      )
  
      # assert after aggregation
      if sk_result.isnan():
          assert result.isnan()
      else:
          assert np.allclose(
              result.numpy(), sk_result, atol=atol
          ), f"result = {result.numpy()}, sk_result = {sk_result}"
  
imports: import os
  import pickle
  import sys
  from functools import partial
  from typing import Callable, Optional
  import numpy
  import pytest
  import torch
  from scipy.stats import entropy
  from torch.distributions.utils import logits_to_probs
  from torch.multiprocessing import Pool, set_start_method
  from torchmetrics import Metric
  from nemo.collections.common.metrics import GlobalAverageLossMetric, Perplexity
label_desc: Use image node_id 7 for calling a global function with example usage: _loss_class_test(rank, worldsize, loss_sum_or_avg, num_measurements, dist_sync_on_step, take_avg_loss, check_dist_sync_on_step, check_batch, atol) without return types