code: def main(cfg: DictConfig):
      root = Path(
          cfg.logging.save_root
      )  # if relative, uses Hydra outputs dir
      model = DVC(**cfg.model)
      logger = WandbLogger(
          save_dir=str(root.absolute()),
          project="DVC",
          config=OmegaConf.to_container(
              cfg
          ),  # saves the Hydra config to wandb
      )
      data = Vimeo90kSeptupletLightning(
          frames_per_group=7,
          **cfg.data,
          pin_memory=cfg.ngpu != 0,
      )
  
      # set up image logging
      rng = np.random.default_rng(cfg.logging.image_seed)
      data.setup()
      val_dataset = data.val_dataset
      log_image_indices = rng.permutation(len(val_dataset))[
          : cfg.logging.num_log_images
      ]
      log_images = torch.stack(
          [val_dataset[ind] for ind in log_image_indices]
      )
      image_logger = WandbImageCallback(log_images)
  
      # run through each stage and optimize
      for stage in sorted(cfg.training_stages.keys()):
          model = run_training_stage(
              stage, root, model, data, logger, image_logger, cfg
          )
  
imports: from pathlib import Path
  import hydra
  import numpy
  import pytorch_lightning
  import torch
  import wandb
  from data_module import Vimeo90kSeptupletLightning
  from dvc_module import DvcModule
  from omegaconf import DictConfig, OmegaConf
  from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
  from pytorch_lightning.loggers import WandbLogger
  from neuralcompression.functional import optical_flow_to_color
  from neuralcompression.models import DVC
label_desc: Use image node_id 3 for calling a global function with example usage: main(cfg) without return types