code: def log_img(self, pl_module, batch, batch_idx, split="train"):
      check_idx = (
          batch_idx if self.log_on_batch_idx else pl_module.global_step
      )
      if (
          self.check_frequency(check_idx)
          and hasattr(
              pl_module, "log_images"
          )  # batch_idx % self.batch_freq == 0
          and callable(pl_module.log_images)
          and
          # batch_idx > 5 and
          self.max_images > 0
      ):
          logger = type(pl_module.logger)
          is_train = pl_module.training
          if is_train:
              pl_module.eval()
  
          gpu_autocast_kwargs = {
              "enabled": self.enable_autocast,  # torch.is_autocast_enabled(),
              "dtype": torch.get_autocast_gpu_dtype(),
              "cache_enabled": torch.is_autocast_cache_enabled(),
          }
          with torch.no_grad(), torch.cuda.amp.autocast(
              **gpu_autocast_kwargs
          ):
              images = pl_module.log_images(
                  batch, split=split, **self.log_images_kwargs
              )
  
          for k in images:
              N = min(images[k].shape[0], self.max_images)
              if not isheatmap(images[k]):
                  images[k] = images[k][:N]
              if isinstance(images[k], torch.Tensor):
                  images[k] = images[k].detach().float().cpu()
                  if self.clamp and not isheatmap(images[k]):
                      images[k] = torch.clamp(images[k], -1.0, 1.0)
  
          self.log_local(
              pl_module.logger.save_dir,
              split,
              images,
              pl_module.global_step,
              pl_module.current_epoch,
              batch_idx,
              pl_module=pl_module
              if isinstance(pl_module.logger, WandbLogger)
              else None,
          )
  
          if is_train:
              pl_module.train()
  
imports: import argparse
  import datetime
  import glob
  import inspect
  import os
  import sys
  from inspect import Parameter
  from typing import Union
  import numpy
  import pytorch_lightning
  import torch
  import torchvision
  import wandb
  from matplotlib import pyplot
  from natsort import natsorted
  from omegaconf import OmegaConf
  from packaging import version
  from PIL import Image
  from pytorch_lightning import seed_everything
  from pytorch_lightning.callbacks import Callback
  from pytorch_lightning.loggers import WandbLogger
  from pytorch_lightning.trainer import Trainer
  from pytorch_lightning.utilities import rank_zero_only
  from sgm.util import exists, instantiate_from_config, isheatmap
label_desc: Use image node_id 3 for calling the ImageLogger obj's underlying member method code with example usage: obj.log_img(pl_module, batch, batch_idx, split) without return types