code: def test_lr_callback_batch_steps_correctly(
      self,
      classifier_module,
      classifier_data,
      policy,
      kwargs,
  ):
      batch_size = 100
      max_epochs = 2
  
      X, y = classifier_data
      num_examples = len(X)
  
      lr_policy = LRScheduler(policy, **kwargs)
      net = NeuralNetClassifier(
          classifier_module(),
          max_epochs=max_epochs,
          batch_size=batch_size,
          callbacks=[lr_policy],
      )
      net.fit(X, y)
  
      total_iterations_per_epoch = num_examples / batch_size
      # 80% of sample used for training by default
      total_training_iterations_per_epoch = (
          0.8 * total_iterations_per_epoch
      )
  
      expected = int(total_training_iterations_per_epoch * max_epochs)
      # pylint: disable=protected-access
      assert lr_policy.batch_idx_ == expected
  
imports: from unittest.mock import Mock
  import numpy
  import pytest
  from sklearn.base import clone
  from torch.optim import SGD
  from torch.optim.lr_scheduler import CosineAnnealingLR
  from torch.optim.lr_scheduler import ExponentialLR
  from torch.optim.lr_scheduler import LambdaLR
  from torch.optim.lr_scheduler import MultiStepLR
  from torch.optim.lr_scheduler import ReduceLROnPlateau
  from torch.optim.lr_scheduler import StepLR
  from torch.optim.lr_scheduler import CyclicLR
  from skorch import NeuralNetClassifier
  from skorch.callbacks.lr_scheduler import WarmRestartLR, LRScheduler
label_desc: Use image node_id 5 for calling the TestLRCallbacks obj's underlying member method code with example usage: obj.test_lr_callback_batch_steps_correctly(classifier_module, classifier_data, policy, kwargs) without return types