class_docstr: Model and training loop for the DVC model.
  
  Combines a pre-defined DVC model with its training loop for use with
  PyTorch Lightning.
  
  Args:
      model: The DVC model to train.
      training_stage: Current stage of training process. One of
          ``("1_motion_estimation", "2_motion_compression",
          "3_motion_compensation", "4_total")``. See DVC paper for details.
      pretrained_model_name: Name of model from CompressAI model zoo to use
          for compressing I-frames.
      pretrained_model_quality_level: Quality level of model from CompressAI.
      num_pframes: Number of P-frames to process for training.
      distortion_type: Type of distortion loss function. Must be from
          ``("MSE")``.
      distortion_lambda: A scaling factor for the distortion term of the
          loss.
      learning_rate: passed to the main network optimizer (i.e. the one that
          adjusts the analysis and synthesis parameters).
      aux_learning_rate: passed to the optimizer that learns the quantiles
          used to build the CDF table for the entropy codder.
      lr_scheduler_params: Used for ``StepLR``, specify ``step_size`` and
          ``gamma``.
class_docstr_tok: ["Model","and","training","loop","for","the","DVC","model",".","Combines","a","pre-defined","DVC","model","with","its","training","loop","for","use","with","PyTorch","Lightning",".","Args",":","model",":","The","DVC","model","to","train",".","training_stage",":","Current","stage","of","training","process",".","One","of","``","(","``","1_motion_estimation","''",",","``","2_motion_compression","''",",","``","3_motion_compensation","''",",","``","4_total","''",")","``",".","See","DVC","paper","for","details",".","pretrained_model_name",":","Name","of","model","from","CompressAI","model","zoo","to","use","for","compressing","I-frames",".","pretrained_model_quality_level",":","Quality","level","of","model","from","CompressAI",".","num_pframes",":","Number","of","P-frames","to","process","for","training",".","distortion_type",":","Type","of","distortion","loss","function",".","Must","be","from","``","(","``","MSE","''",")","``",".","distortion_lambda",":","A","scaling","factor","for","the","distortion","term","of","the","loss",".","learning_rate",":","passed","to","the","main","network","optimizer","(","i.e",".","the","one","that","adjusts","the","analysis","and","synthesis","parameters",")",".","aux_learning_rate",":","passed","to","the","optimizer","that","learns","the","quantiles","used","to","build","the","CDF","table","for","the","entropy","codder",".","lr_scheduler_params",":","Used","for","``","StepLR","``",",","specify","``","step_size","``","and","``","gamma","``","."]
docstr: Used by trainer to rebuild the model after a training run.
docstr_tok: ["Used","by","trainer","to","rebuild","the","model","after","a","training","run","."]
code: def recompose_model(self, model: DVC):
      """Used by trainer to rebuild the model after a training run."""
      return self.model.recompose_model(model)
  
imports: from typing import NamedTuple, Optional, Tuple
  import torch
  import torch.optim
  import torchmetrics
  from _utils import DvcStage, DvcStage1, DvcStage2, DvcStage3, DvcStage4and5, LossFunctions
  from compressai.zoo import models
  from pytorch_lightning import LightningModule
  from torch import Tensor
  import neuralcompression.functional
  from neuralcompression.models import DVC
label_desc: Use image node_id 2 for calling the DvcModule obj's underlying member method code with example usage: obj.recompose_model(model) and returns: self