class_docstr: Encapsulates model + loss + optimizer + metrics in a PL module
class_docstr_tok: ["Encapsulates","model","+","loss","+","optimizer","+","metrics","in","a","PL","module"]
code: def training_step(
      self, batch, batch_idx: int, optimizer_idx: Optional[int] = None
  ) -> Tensor:
      recon_frames, rate_args = self(batch)  # [B, T, 3, H, W], [Tensor]
  
      dist_loss_frame = F.mse_loss(
          recon_frames, batch.video_tensor, reduction="none"
      ).mean(dim=(0, 2, 3, 4))
      rate_loss_frame = self.rate_loss(
          batch, rate_args, per_frame=True
      )  # [T]
      weights = 10.0 * torch.ones_like(rate_loss_frame)
      weights[0] = 1.0
      combined_loss_frame = (
          self._distortion_loss_scaling * dist_loss_frame
          + rate_loss_frame
      )
      combined_loss_weighted = (combined_loss_frame * weights).mean()
  
      with torch.no_grad():
          psnr = peak_signal_noise_ratio(
              recon_frames,
              batch.video_tensor,
              data_range=1.0,
              dim=(2, 3, 4),
          ).item()
          psnr_int = peak_signal_noise_ratio(
              (recon_frames * 255 + 0.5).to(torch.uint8),
              (batch.video_tensor * 255 + 0.5).to(torch.uint8),
              data_range=255,
              dim=(2, 3, 4),
          ).item()
      self.log_dict(
          {
              "distortion_loss": dist_loss_frame.mean(),
              "rate_loss": rate_loss_frame.mean(),
              "combined_loss": combined_loss_frame.mean(),
              "combined_loss_weighted": combined_loss_weighted,
              "PSNR": psnr,
              "PSNR_int": psnr_int,
          },
          sync_dist=True,
      )
      return combined_loss_weighted
  
imports: from typing import Any, List, Optional, Tuple, Union
  import hydra
  import pytorch_lightning
  import torch
  import torch.nn.functional
  from datamodules.video_data_api import VideoData
  from model_pipeline import VCTPipeline
  from omegaconf import DictConfig
  from torch import Tensor
  from torchmetrics.functional import peak_signal_noise_ratio
  from utils.hydra_tools import OmegaConf
label_desc: Use image node_id 4 for calling the VCTModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx, optimizer_idx) and returns: combined_loss_weighted