code: def _class_test(
      rank: int,
      worldsize: int,
      preds: torch.Tensor,
      target: torch.Tensor,
      metric_class: Metric,
      sk_metric: Callable,
      dist_sync_on_step: bool,
      metric_args: dict = {},
      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
          preds: torch tensor with predictions
          target: torch tensor with targets
          metric_class: lightning metric class that should be tested
          sk_metric: callable function that is used for comparison
          dist_sync_on_step: bool, if true will synchronize metric state across
              processes at each ``forward()``
          metric_args: 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)
      """
      # Instanciate lightning metric
      metric = metric_class(
          dist_sync_on_step=dist_sync_on_step, **metric_args
      )
  
      # verify metrics work after being loaded from pickled state
      pickled_metric = pickle.dumps(metric)
      metric = pickle.loads(pickled_metric)
  
      for i in range(rank, NUM_BATCHES, worldsize):
          batch_result = metric(preds[i], target[i])
  
          if metric.dist_sync_on_step:
              if rank == 0:
                  ddp_preds = torch.stack(
                      [preds[i + r] for r in range(worldsize)]
                  )
                  ddp_target = torch.stack(
                      [target[i + r] for r in range(worldsize)]
                  )
                  sk_batch_result = sk_metric(ddp_preds, ddp_target)
                  # assert for dist_sync_on_step
                  if check_dist_sync_on_step:
                      assert np.allclose(
                          batch_result.numpy(),
                          sk_batch_result,
                          atol=atol,
                      )
          else:
              sk_batch_result = sk_metric(preds[i], target[i])
              # assert for batch
              if check_batch:
                  assert np.allclose(
                      batch_result.numpy(), sk_batch_result, atol=atol
                  )
  
      # check on all batches on all ranks
      result = metric.compute()
      assert isinstance(result, torch.Tensor)
  
      total_preds = torch.stack([preds[i] for i in range(NUM_BATCHES)])
      total_target = torch.stack(
          [target[i] for i in range(NUM_BATCHES)]
      )
      sk_result = sk_metric(total_preds, total_target)
  
      # assert after aggregation
      assert np.allclose(result.numpy(), sk_result, atol=atol)
  
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 2 for calling a global function with example usage: _class_test(rank, worldsize, preds, target, metric_class, sk_metric, dist_sync_on_step, metric_args, check_dist_sync_on_step, check_batch, atol) without return types