code: def on_train_epoch_end(
      self, trainer: Trainer, pl_module: LightningModule
  ) -> None:
      if not _TORCHVISION_AVAILABLE:
          return
  
      images, _ = next(
          iter(
              DataLoader(
                  trainer.datamodule.mnist_val,
                  batch_size=self.num_samples,
              )
          )
      )
      images_flattened = images.view(images.size(0), -1)
  
      # generate images
      with torch.no_grad():
          pl_module.eval()
          images_generated = pl_module(
              images_flattened.to(pl_module.device)
          )
          pl_module.train()
  
      if trainer.current_epoch == 0:
          save_image(
              self._to_grid(images),
              f"grid_ori_{trainer.current_epoch}.png",
          )
      save_image(
          self._to_grid(images_generated.reshape(images.shape)),
          f"grid_generated_{trainer.current_epoch}.png",
      )
  
imports: from os import path
  from typing import Optional, Tuple
  import torch
  import torch.nn.functional
  from lightning.pytorch import LightningDataModule, LightningModule, Trainer, callbacks, cli_lightning_logo
  from lightning.pytorch.cli import LightningCLI
  from lightning.pytorch.demos.mnist_datamodule import MNIST
  from lightning.pytorch.utilities import rank_zero_only
  from lightning.pytorch.utilities.imports import _TORCHVISION_AVAILABLE
  from torch import nn
  from torch.utils.data import DataLoader, random_split
label_desc: Use image node_id 3 for calling the ImageSampler obj's underlying member method code with example usage: obj.on_train_epoch_end(trainer, pl_module) without return types