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 log_images(self, trainer, base_key, image_dict, global_step):
      # check if we have images to log
      # if we do, then concatenate time along x-axis and batch along y-axis
      # and write
      keys = ("flow", "image1", "image2", "image2_est")
      for key in keys:
          if image_dict.get(key) is not None:
              caption = f"{key} (y-axis: batch, x-axis: time)"
              mosaic = torch.cat(image_dict[key], dim=-1)
              mosaic = torch.cat(list(mosaic), dim=-2)
              if key == "flow":
                  mosaic = optical_flow_to_color(mosaic.unsqueeze(0))[0]
              mosaic = torch.clip(mosaic, min=0, max=1.0)
              trainer.logger.experiment.log(
                  {
                      f"{base_key}/{key}": wandb.Image(
                          mosaic, caption=caption
                      ),
                      "global_step": 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 3 for calling the WandbImageCallback obj's underlying member method code with example usage: obj.log_images(trainer, base_key, image_dict, global_step) without return types