code: def run_training_stage(
      stage, root, model, data, logger, image_logger, cfg
  ):
      """Run a single training stage based on the stage config."""
      print(f"training stage: {stage}")
      stage_cfg = cfg.training_stages[stage]
      if stage_cfg.save_dir is None:
          save_dir = root / stage
      else:
          save_dir = Path(stage_cfg.save_dir)
  
      if (
          not cfg.checkpoint.overwrite
          and not cfg.checkpoint.resume_training
          and len(list(save_dir.glob("*.ckpt"))) > 0
      ):
          raise RuntimeError(
              "Checkpoints detected in save directory: set resume_training=True "
              "to restore trainer state from these checkpoints, "
              "or set overwrite=True to ignore them."
          )
  
      save_dir.mkdir(exist_ok=True, parents=True)
      last_checkpoint = save_dir / "last.ckpt"
      if (
          not last_checkpoint.exists()
          or cfg.checkpoint.overwrite is True
      ):
          last_checkpoint = None
  
      lightning_model = DvcModule(
          model, **merge_configs(cfg.module, stage_cfg.module)
      )
  
      trainer = pl.Trainer(
          **merge_configs(cfg.trainer, stage_cfg.trainer),
          logger=logger,
          callbacks=[
              LearningRateMonitor(),
              ModelCheckpoint(
                  dirpath=save_dir, **cfg.checkpoint.model_checkpoint
              ),
              image_logger,
          ],
          resume_from_checkpoint=last_checkpoint,
      )
  
      trainer.fit(lightning_model, datamodule=data)
  
      return lightning_model.recompose_model(model)
  
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 2 for calling a global function with example usage: run_training_stage(stage, root, model, data, logger, image_logger, cfg) and returns: lightning_model