code: def main(rank, world_size, args):
      start = time.time()
  
      if world_size > 1:
          torch.distributed.init_process_group(
              backend="gloo",
              init_method="env://",
              world_size=world_size,
              rank=rank,
          )
          torch.cuda.set_device(rank % torch.cuda.device_count())
  
      raw_args = args
  
      args = convert_namespace_to_omegaconf(args)
      if args.common.seed is not None:
          np.random.seed(args.common.seed)
          utils.set_torch_seed(args.common.seed)
  
      (
          models,
          model_args,
          task,
      ) = checkpoint_utils.load_model_ensemble_and_task(
          [raw_args.path], arg_overrides={"data": args.task.data}
      )
  
      tgt_dict = task.target_dictionary
  
      for model in models:
          model.prepare_for_inference_(args)
          model.cuda().eval()
          if raw_args.fp16:
              model = model.half()
      model = models[0]
  
      config = ExpressiveCodeDataConfig(args.task.data)
  
      dataset = CodeDataset(
          manifest=config.manifests[raw_args.eval_subset],
          dictionary=task.source_dictionary,
          dur_dictionary=task.source_duration_dictionary,
          f0_dictionary=task.source_f0_dictionary,
          config=config,
          discrete_dur=task.cfg.discrete_duration,
          discrete_f0=task.cfg.discrete_f0,
          log_f0=task.cfg.log_f0,
          normalize_f0_mean=task.cfg.normalize_f0_mean,
          normalize_f0_std=task.cfg.normalize_f0_std,
          interpolate_f0=task.cfg.interpolate_f0,
          shifts=task.cfg.stream_shifts,
          return_filename=True,
          strip_filename=False,
          return_continuous_f0=raw_args.dequantize_prosody,
      )
  
      if raw_args.filter_names:
          dataset = FilterNamesDataset(dataset, raw_args.filter_names)
  
      criterion = task.build_criterion(model_args.criterion)
  
      name2metric = {
          "continuation": continuation,
          "teacher_force_everything": teacher_force_everything,
          "correlation": correlation,
      }
  
      name2keys = {
          "continuation": (
              "Token BLEU3",
              "Duration NLL",
              "Duration MAE",
              "F0 NLL",
              "F0 MAE",
              "F0 sum",
              "F0 sum_sq",
              "Dur sum",
              "Dur sum_sq",
          ),
          "teacher_force_everything": (
              "token_loss",
              "duration_loss",
              "f0_loss",
          ),
          "correlation": ("Duration corr", "F0 corr"),
      }
      metric_name = raw_args.metric
  
      metric = name2metric[metric_name]
      results = metric(
          raw_args,
          dataset,
          model,
          criterion,
          tgt_dict,
          rank,
          world_size,
      )
  
      values = None
  
      if metric_name not in [
          "correlation",
      ]:
          values, normalizers = results
          values = maybe_aggregate_normalize(
              values, normalizers, world_size
          )
      elif metric_name == "correlation":
          values = maybe_aggregate_correlations(results, world_size)
      else:
          assert False
  
      assert values is not None
      summary = dict(zip(name2keys[raw_args.metric], values.tolist()))
      if metric_name == "continuation":
          summary["F0 Std"] = np.sqrt(
              -summary["F0 sum"] ** 2 + summary["F0 sum_sq"]
          )
          summary["Dur Std"] = np.sqrt(
              -summary["Dur sum"] ** 2 + summary["Dur sum_sq"]
          )
          del summary["F0 sum"]
          del summary["F0 sum_sq"]
          del summary["Dur sum"]
          del summary["Dur sum_sq"]
  
      summary["metric"] = metric_name
  
      if rank == 0:
          print(summary)
          if raw_args.wandb:
              wandb_results(summary, raw_args)
          print("# finished in ", time.time() - start, "seconds")
  
imports: import os
  import numpy
  import scipy
  import torch
  import torch.multiprocessing
  from fairseq import checkpoint_utils, options
  from fairseq.data.codedataset import CodeDataset, ExpressiveCodeDataConfig
  from fairseq.dataclass.utils import convert_namespace_to_omegaconf
  from torch.utils.data import DataLoader, DistributedSampler
  from fairseq.utils import move_to_cuda
  from fairseq import utils
  from fairseq.criterions.speech_ulm_criterion import nll_loss, mae_loss
  import time
  from types import SimpleNamespace
  import sys, pathlib
  from naive_decoder import Naive_F0_Decoder
  from inference_dataset import InferenceDataset, explode_batch
  from sample.sample import do_sampling, TemperatureDecoder, FilterNamesDataset
label_desc: Use image node_id 5 for calling a global function with example usage: main(rank, world_size, args) without return types