class_docstr: Logs the input and output images of a module.
  
  Images are stacked into a mosaic, with different clips from top to bottom
  and time progressing from left to right.
  
  Args:
      batch: A set of images to log at the end of each validation epoch.
class_docstr_tok: ["Logs","the","input","and","output","images","of","a","module",".","Images","are","stacked","into","a","mosaic",",","with","different","clips","from","top","to","bottom","and","time","progressing","from","left","to","right",".","Args",":","batch",":","A","set","of","images","to","log","at","the","end","of","each","validation","epoch","."]
code: def on_validation_end(self, trainer, pl_module):
      image_dict = {}
      batch = self.batch.to(
          device=pl_module.device, dtype=pl_module.dtype
      )
      batch, _ = pl_module.compress_iframe(batch)  # bpp_total w/o grads
      image1 = batch[:, 0]
      for i in range(pl_module.num_pframes):
          image2 = batch[:, i + 1]
          _, images = pl_module.model.compute_batch_loss(image1, image2)
          image1 = images.image2_est  # images are detached
  
          image_dict = self.append_images(image_dict, images)
  
      self.log_images(
          trainer,
          f"log_images_stage_{pl_module.training_stage}",
          image_dict,
          pl_module.global_step,
      )
  
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 4 for calling the WandbImageCallback obj's underlying member method code with example usage: obj.on_validation_end(trainer, pl_module) without return types