code: def on_train_batch_end(  # type: ignore[override]
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      outputs: STEP_OUTPUT,
      batch: Any,
      batch_idx: int,
      dataloader_idx: int,
  ) -> None:
      if trainer.max_epochs is None:
          raise ValueError(
              "trainer.max_epochs is none, cannot get training progress"
          )
      else:
          elapsed_duration = float(trainer.current_epoch) / float(
              trainer.max_epochs
          )
      if self.__should_selective_backprop(
          elapsed_duration,
          batch_idx,
          self.start,
          self.end,
          self.interrupt,
      ):
          current_batch_size = len(batch[1])
          ideal_batch_size = int(self.keep * self.batch_size)
          assert (
              current_batch_size == ideal_batch_size
          ), "Batch size is not right."
  
imports: import os
  import torch
  import torchmetrics
  from torch import nn
  from torchvision.datasets import CIFAR10
  from torch.utils.data import DataLoader
  import torchvision.transforms
  from torch.utils.data.sampler import SequentialSampler
  import pytorch_lightning
  from pytorch_lightning.callbacks import Callback
  from pytorch_lightning.utilities.types import STEP_OUTPUT
  from typing import Any
  from bigdl.nano.pytorch.lightning import LightningModule
  from bigdl.nano.pytorch import Trainer
  from bigdl.nano.pytorch.vision.models import vision
  from bigdl.nano.pytorch.algorithms.selective_backprop import SelectiveBackprop
label_desc: Use image node_id 3 for calling the CheckBatchSize obj's underlying member method code with example usage: obj.on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx) without return types