code: def on_train_batch_start(
      self,
      trainer: "pl.Trainer",
      pl_module: "pl.LightningModule",
      batch,
      batch_idx,
      unused: int = 0,
  ) -> None:
      if self.warmup_on_epoch:
          return
      if batch_idx > self.warmup_epochs or pl_module.current_epoch > 0:
          for lr, opt in zip(self.lrs, pl_module.optimizers()):
              assert opt.param_groups[0]["lr"] == lr * self.world_size
      else:
          for lr, opt in zip(self.lrs, pl_module.optimizers()):
              diff = (
                  (1.0 - 1.0 / self.world_size)
                  * batch_idx
                  / self.warmup_epochs
                  + 1.0 / self.world_size
              )
              assert (
                  opt.param_groups[0]["lr"]
                  == lr * self.world_size * diff
              )
  
imports: import os
  import pytest
  from unittest import TestCase
  from test.pytorch.utils._train_torch_lightning import create_data_loader, data_transform
  from test.pytorch.utils._train_torch_lightning import create_test_data_loader
  from bigdl.nano.pytorch.vision.models import vision
  from bigdl.nano.pytorch import Trainer
  from bigdl.nano.utils.pytorch import TORCH_VERSION_LESS_1_10, TORCH_VERSION_LESS_2_0
  import pytorch_lightning
  import torch
  from torch import nn
  import torch.nn.functional
label_desc: Use image node_id 2 for calling the CheckWarmupCallback obj's underlying member method code with example usage: obj.on_train_batch_start(trainer, pl_module, batch, batch_idx, unused) without return types