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 test_step(
      self,
      batch: VideoData,
      batch_idx: int,
      run_fwd: bool = True,
  ):
      B, T, C, H, W = batch.shape
      assert (
          B == 1
      ), "Metrics calculation only supported for batch size 1."
      assert self.training is False
  
      with torch.no_grad():
          video_ref_uint8 = (batch.video_tensor[0] * 255 + 0.5).to(
              torch.uint8
          )
          if run_fwd:
              self.model._code_to_strings = False
              forward_pass = self.model(batch)
              recon, rate_args = forward_pass[-2], forward_pass[-1][0]
  
              # take the single batch and clamp it:
              recon_clamped = torch.clamp(recon[0], 0.0, 1.0)
              # caveat: real compression people use YCbCr
              recon_uint8 = (recon_clamped * 255 + 0.5).to(torch.uint8)
              cts_psnr = peak_signal_noise_ratio(
                  recon_uint8, video_ref_uint8, data_range=255
              ).item()
              cts_rate = self.model.compute_rate(rate_args).item() / (
                  H * W * T
              )
              # Compute the rate excluding the first frame since we are technically
              # not compressing it. Generally even if the compression is not great
              # for that first frame, its coding cost will amportize over the video
              # sequence length
              cts_rate_x1 = self.model.compute_rate(
                  rate_args[1:]
              ).item() / (H * W * (T - 1))
              self.log_dict(
                  {
                      f"cts_psnr": cts_psnr,
                      f"cts_rate": cts_rate,
                      f"cts_rate_x1": cts_rate_x1,
                  },
                  sync_dist=True,
              )
  
          *output, bits_per_frame = self.model.compress_video(
              batch, force_cpu=False
          )
          recon = self.model.decompress_video(
              batch.shape, bottleneck_args=output, force_cpu=False
          )
          # clamp the reconstructed video (it's batch size 1, so pick [0])
          recon_clamped = torch.clamp(recon[0], 0.0, 1.0)
  
          # PSNR
          # caveat: real compression people use YCbCr
          recon_uint8 = (recon_clamped * 255 + 0.5).to(torch.uint8)
          psnr = peak_signal_noise_ratio(
              recon_uint8, video_ref_uint8, data_range=255
          ).item()
          # bits per pixel
          bpp = sum(bits_per_frame) / (H * W * T)
          # Compute the rate excluding the first frame since we are technically
          # not compressing it. Generally even if the compression is not great
          # for that first frame, its coding cost will amportize over the video
          # sequence length
          bpp_xfirst = sum(bits_per_frame[1:]) / (H * W * (T - 1))
          print(psnr, bpp, bpp_xfirst)
          self.log_dict(
              {
                  f"psnt_int": psnr,
                  f"bpp": bpp,
                  f"bpp_xfirst": bpp_xfirst,
              },
              sync_dist=True,
          )
      return psnr, bpp, bpp_xfirst
  
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 6 for calling the VCTModule obj's underlying member method code with example usage: obj.test_step(batch, batch_idx, run_fwd) and returns: psnr,  bpp,  bpp_xfirst